using System;
using System.Collections.Generic;
using UAF;
using UAF.Interfaces.Logging;
using System.IO;
using System.Runtime.InteropServices;
using System.Xml.Linq;
using System.Diagnostics;

namespace DataFileAnalyzer
{
    /// <summary>
    /// Stores informations associated to the File Reference Numbers of a folder of a NTFS volume.
    /// </summary>
    class NtfsFolderDictionary
    {
        private ILog m_log = new NullLog();

        private string m_folder;
        private UInt64 m_folderFRN;
        private string m_drive;
        private List<string> m_excludeFolders = new List<string>();
        private List<UInt64> m_excludedFRNs = new List<UInt64>();
        private IntPtr m_usnJournalRootHandle = IntPtr.Zero;

        /// <summary>
        /// Informations about a FRN
        /// </summary>
        private class FRNInfo
        {
            public string Name { get; set; }
            public UInt64 ParentFRN { get; set; }
            public string NormalizedFullPath { get; set; }
            public bool IsFolder { get; set; }

            public FRNInfo(string _name, UInt64 _parentFrn, bool _isFolder)
            {
                Name = _name;
                ParentFRN = _parentFrn;
                NormalizedFullPath = null;
                IsFolder = _isFolder;
            }
        }
        private Dictionary<UInt64, FRNInfo> m_frnInfos = new Dictionary<UInt64, FRNInfo>();

        /// <summary>
        /// Constructor
        /// </summary>
        public NtfsFolderDictionary(ILog _log)
        {
            m_log = _log;

            m_folder = null;
            m_drive = null;
            m_excludeFolders = null;
            m_usnJournalRootHandle = IntPtr.Zero;
        }

        /// <summary>
        /// The data folder of the library has changed
        /// </summary>
        public void OnDataFolderChanged(string _folder, List<string> _excludedFolders)
        {
            m_folder = _folder;
            m_drive = Path.GetPathRoot(_folder);
            m_excludeFolders = new List<string>(_excludedFolders);
        }

        /// <summary>
        /// Set the journal handle
        /// </summary>
        public void SetJournalHandle(IntPtr _journalHandle)
        {
            m_usnJournalRootHandle = _journalHandle;
        }

        /// <summary>
        /// Does _normalizedPath points to a file in this folder?
        /// </summary>
        private bool IsPathInFolder(string _normalizedPath)
        {
            if (_normalizedPath.StartsWith(m_folder))
            {
                foreach (string folder in m_excludeFolders)
                {
                    if (_normalizedPath.StartsWith(folder))
                        return false;
                }
                return true;
            }
            return false;
        }

        /// <summary>
        /// Get all the files of the folder
        /// </summary>
        public IEnumerable<string> GetAllFiles()
        {
            foreach (KeyValuePair<UInt64, FRNInfo> pair in m_frnInfos)
            {
                FRNInfo info = pair.Value;
                if (!info.IsFolder)
                    yield return info.NormalizedFullPath;
            }
        }

        /// <summary>
        /// Get the normalized file path of a FRN
        /// </summary>
        public string GetFRNFullPath(UInt64 _frn)
        {
            FRNInfo info;
            if (m_frnInfos.TryGetValue(_frn, out info))
                return info.NormalizedFullPath;
            return null;
        }

        /// <summary>
        /// Is this FRN in the dictionary?
        /// </summary>
        public bool IsFRNInDictionary(UInt64 _frn)
        {
            return m_frnInfos.ContainsKey(_frn);
        }

        /// <summary>
        /// Is _childFrn a child of _parentFrn?
        /// </summary>
        private bool IsChildOf(UInt64 _childFrn, UInt64 _parentFrn)
        {
            FRNInfo info;
            while (m_frnInfos.TryGetValue(_childFrn, out info))
            {
                if (info.ParentFRN == _parentFrn)
                    return true;
                _childFrn = info.ParentFRN;
            }
            return false;
        }

        /// <summary>
        /// Get the children of an element of the dictionary
        /// </summary>
        private void GetChildrenOf(UInt64 _frn, out List<KeyValuePair<UInt64, FRNInfo>> _children)
        {
            _children = new List<KeyValuePair<UInt64, FRNInfo>>();
            foreach (KeyValuePair<UInt64, FRNInfo> pair in m_frnInfos)
            {
                if (pair.Value.ParentFRN == _frn || IsChildOf(pair.Key, _frn))
                    _children.Add(pair);
            }
        }

        /// <summary>
        /// Is _frn a child of _rootFrn given all the FRN infos _allInfos?
        /// </summary>
        private static bool IsFRNInFolder(UInt64 _frn, UInt64 _rootFrn, List<UInt64> _excludedFRNs, Dictionary<UInt64, FRNInfo> _allInfos)
        {
            if (_frn == _rootFrn)
                return true;
            if (_excludedFRNs.Contains(_frn))
                return false;

            FRNInfo info = null;
            while (_allInfos.TryGetValue(_frn, out info))
            {
                if (info.ParentFRN == _rootFrn)
                    return true;
                if (_excludedFRNs.Contains(info.ParentFRN))
                    return false;
                _frn = info.ParentFRN;
            }

            return false;
        }

        /// <summary>
        /// Parse the root folder to initialize the dictionary
        /// </summary>
        public bool InitFromFolder()
        {
            m_frnInfos.Clear();
            m_excludedFRNs.Clear();

            DateTime startTime = DateTime.UtcNow;

            if (!Win32Api.GetFileReferenceFromPath(m_folder, out m_folderFRN))
            {
                m_log.Error("NtfsVolumeDictionary could not initialize!");
                return false;
            }
            UInt64 excludedFRN;
            foreach (string excludedFolder in m_excludeFolders)
            {
                if (!Win32Api.GetFileReferenceFromPath(excludedFolder, out excludedFRN))
                {
                    m_log.Error("NtfsVolumeDictionary could not initialize!");
                    return false;
                }
                m_excludedFRNs.Add(excludedFRN);
            }

            List<string> addedFiles;
            AddFolder(null, 0, m_folderFRN, out addedFiles);

            m_log.Message("NtfsVolumeDictionary: Init done in {0}. {1} files and folders in data folder.",
                (DateTime.UtcNow - startTime).ToReadableString(),
                m_frnInfos.Count);

            return true;
        }

        /// <summary>
        /// DEPRECATED: We can miss files in the folder if the journal has been (re)created
        /// since the folder was created.
        /// Read the MFT entrie before _highUsn to init the dictionary.
        /// </summary>
        public bool InitFromMFT(Int64 _highUsn, out Int64 _lastUsnRead)
        {
            _lastUsnRead = 0;
            m_frnInfos.Clear();
            m_excludedFRNs.Clear();

            DateTime startTime = DateTime.UtcNow;

            if (!Win32Api.GetFileReferenceFromPath(m_folder, out m_folderFRN))
            {
                m_log.Error("NtfsVolumeDictionary could not initialize!");
                return false;
            }
            UInt64 excludedFRN;
            foreach (string excludedFolder in m_excludeFolders)
            {
                if (!Win32Api.GetFileReferenceFromPath(excludedFolder, out excludedFRN))
                {
                    m_log.Error("NtfsVolumeDictionary could not initialize!");
                    return false;
                }
                m_excludedFRNs.Add(excludedFRN);
            }

            Win32Api.MFT_ENUM_DATA enumData;
            enumData.StartFileReferenceNumber = 0;
            enumData.LowUsn = 0;
            enumData.HighUsn = _highUsn;
            Int32 enumDataSize = Marshal.SizeOf(enumData);
            IntPtr enumDataBuffer = Marshal.AllocHGlobal(enumDataSize);
            Win32Api.ZeroMemory(enumDataBuffer, enumDataSize);
            Marshal.StructureToPtr(enumData, enumDataBuffer, true);

            int bufferSize = sizeof(UInt64) + 10000;
            IntPtr bufferPointer = Marshal.AllocHGlobal(bufferSize);
            Win32Api.ZeroMemory(bufferPointer, bufferSize);
            uint bytesRead = 0;

            Dictionary<UInt64, FRNInfo> allFRNs = new Dictionary<UInt64, FRNInfo>();
            while (Win32Api.DeviceIoControl(
                m_usnJournalRootHandle,
                Win32Api.FSCTL_ENUM_USN_DATA,
                enumDataBuffer,
                enumDataSize,
                bufferPointer,
                bufferSize,
                out bytesRead,
                IntPtr.Zero))
            {
                // Skip the next USN stored at the beginning of the buffer
                bytesRead -= sizeof(Int64);

                IntPtr usnRecordPointer = IntPtr.Add(bufferPointer, sizeof(Int64));
                while (bytesRead > 0)
                {
                    Win32Api.UsnEntry usnEntry = new Win32Api.UsnEntry(usnRecordPointer);
                    FRNInfo frnInfo = new FRNInfo(usnEntry.Name.ToLower(), usnEntry.ParentFileReferenceNumber, usnEntry.IsFolder);
                    allFRNs.Add(usnEntry.FileReferenceNumber, frnInfo);
                    if (usnEntry.USN > _lastUsnRead)
                        _lastUsnRead = usnEntry.USN;

                    usnRecordPointer = new IntPtr(usnRecordPointer.ToInt32() + usnEntry.RecordLength);
                    usnRecordPointer = IntPtr.Add(usnRecordPointer, (int)usnEntry.RecordLength);
                    bytesRead -= usnEntry.RecordLength;
                }
                Marshal.WriteInt64(enumDataBuffer, Marshal.ReadInt64(bufferPointer, 0));
            }

            Marshal.FreeHGlobal(enumDataBuffer);
            Marshal.FreeHGlobal(bufferPointer);

            DateTime endMFTParseTime = DateTime.UtcNow;

            // Add the FRN infos to the dictionary if they are in the right folder
            foreach (KeyValuePair<UInt64, FRNInfo> pair in allFRNs)
            {
                FRNInfo frnInfo = pair.Value;
                if (IsFRNInFolder(pair.Key, m_folderFRN, m_excludedFRNs, allFRNs))
                    m_frnInfos.Add(pair.Key, frnInfo);
            }

            // Compute the full paths
            foreach (KeyValuePair<UInt64, FRNInfo> pair in m_frnInfos)
                ComputeFullPath(pair.Key, pair.Value);

            DateTime utcNow = DateTime.UtcNow;
            m_log.Message("NtfsVolumeDictionary: Init done in {0}. MFT parsing done in {3}. {1} files parsed. {2} files in data folder.",
                (utcNow - startTime).ToReadableString(),
                allFRNs.Count,
                m_frnInfos.Count,
                (utcNow - endMFTParseTime).ToReadableString());

            return true;
        }

        /// <summary>
        /// Build the full path of a FRN info
        /// </summary>
        private void ComputeFullPath(UInt64 _frn, FRNInfo _info)
        {
            if (_frn != m_folderFRN)
            {
                string fullPath = _info.Name;
                FRNInfo parentInfo = _info;
                while (parentInfo.ParentFRN != m_folderFRN && m_frnInfos.TryGetValue(parentInfo.ParentFRN, out parentInfo))
                    fullPath = parentInfo.Name + '/' + fullPath;
                _info.NormalizedFullPath = m_folder + '/' + fullPath;
            }
            else
            {
                _info.NormalizedFullPath = m_folder;
            }
        }

        /// <summary>
        /// Add a file to the dictionary
        /// </summary>
        public void AddFile(
            string _name,
            UInt64 _frn,
            UInt64 _parentFrn)
        {
            FRNInfo newInfo = new FRNInfo(_name.ToLower(), _parentFrn, false);
            m_frnInfos[_frn] = newInfo;
            ComputeFullPath(_frn, newInfo);
        }

        /// <summary>
        /// Remove a file from the dictionary
        /// </summary>
        public void RemoveFile(UInt64 _frn)
        {
            m_frnInfos.Remove(_frn);
        }

        /// <summary>
        /// Rename or move file
        /// </summary>
        public void MoveFile(UInt64 _frn, string _newName, UInt64 _newParentFrn)
        {
            FRNInfo info;
            if (m_frnInfos.TryGetValue(_frn, out info))
            {
                info.Name = _newName.ToLower();
                info.ParentFRN = _newParentFrn;
                ComputeFullPath(_frn, info);
            }
        }


        /// <summary>
        /// Add a folder to the dictionary. All the sub-folders and files are added.
        /// _addedFiles is filled with the full paths of the files added to the dictionary.
        /// WARNING: Rough version. Could be optimized.
        /// </summary>
        public void AddFolder(string _name, UInt64 _frn, UInt64 _parentFrn, out List<string> _addedFiles)
        {
            _addedFiles = new List<string>();
            UInt64 frn;
            FRNInfo newInfo;
            Dictionary<string, UInt64> folderPathToFRN = new Dictionary<string, UInt64>();
            string folderFullPath;

            if (string.IsNullOrEmpty(_name) && _parentFrn == m_folderFRN)
            {
                // We're adding the root folder
                FRNInfo rootInfo = new FRNInfo(Path.GetFileName(m_folder), 0, true);
                rootInfo.NormalizedFullPath = m_folder;
                m_frnInfos[_parentFrn] = rootInfo;

                folderFullPath = m_folder;
                folderPathToFRN.Add(m_folder, m_folderFRN);
            }
            else
            {
                FRNInfo parentInfo;
                if (!m_frnInfos.TryGetValue(_parentFrn, out parentInfo))
                    return;

                folderFullPath = parentInfo.NormalizedFullPath + '/' + _name.ToLower();

                // Check that this is not an excluded folder
                if (!IsPathInFolder(folderFullPath))
                    return;

                folderPathToFRN.Add(parentInfo.NormalizedFullPath, _parentFrn);
            }

            if (Win32Api.GetFileReferenceFromPath(folderFullPath, out frn))
            {
                BasicFileEnumerator enumerator = new BasicFileEnumerator();
                foreach (string path in enumerator.GetAllFolders(folderFullPath, m_excludeFolders))
                {
                    if (Win32Api.GetFileReferenceFromPath(path, out frn))
                    {
                        string fileName = Path.GetFileName(path);
                        string parentPath = PathUtils.Normalize(Path.GetDirectoryName(path));
                        UInt64 parentFrn;
                        if (folderPathToFRN.TryGetValue(parentPath, out parentFrn))
                        {
                            newInfo = new FRNInfo(fileName.ToLower(), parentFrn, true);
                            newInfo.NormalizedFullPath = path;
                            m_frnInfos[frn] = newInfo;
                            folderPathToFRN.Add(path, frn);
                        }
                    }
                }
                foreach (string path in enumerator.GetAllFiles(folderFullPath, m_excludeFolders))
                {
                    if (Win32Api.GetFileReferenceFromPath(path, out frn))
                    {
                        string fileName = Path.GetFileName(path);
                        string parentPath = PathUtils.Normalize(Path.GetDirectoryName(path));
                        UInt64 parentFrn;
                        if (folderPathToFRN.TryGetValue(parentPath, out parentFrn))
                        {
                            newInfo = new FRNInfo(fileName.ToLower(), parentFrn, false);
                            newInfo.NormalizedFullPath = PathUtils.Normalize(path);
                            m_frnInfos[frn] = newInfo;
                            _addedFiles.Add(newInfo.NormalizedFullPath);
                        }
                    }
                }
            }
        }

        /// <summary>
        /// Remove a folder from the dictionary. All the sub-folders and files
        /// will be removed. Add all the removed files to _removedFiles.
        /// WARNING: Rough version. Could be optimized.
        /// </summary>
        public void RemoveFolder(UInt64 _frn, out List<string> _removedFiles)
        {
            _removedFiles = new List<string>();
            FRNInfo info;
            if (m_frnInfos.TryGetValue(_frn, out info) && info.IsFolder)
            {
                List<KeyValuePair<UInt64, FRNInfo>> removedList;
                GetChildrenOf(_frn, out removedList);

                m_frnInfos.Remove(_frn);
                foreach (KeyValuePair<UInt64, FRNInfo> pair in removedList)
                {
                    FRNInfo childInfo = pair.Value;
                    if (!childInfo.IsFolder)
                        _removedFiles.Add(childInfo.NormalizedFullPath);

                    m_frnInfos.Remove(pair.Key);
                }
            }
        }

        /// <summary>
        /// Move a folder inside the dictionary. All the sub-folders and files
        /// will be updated accordingly. Add the old path of all the moved files
        /// to _removedFiles and add their new path to _addedFiles.
        /// WARNING: Rough version. Could be optimized.
        /// </summary>
        public void MoveFolder(UInt64 _frn, string _newName, UInt64 _newParentFrn, out List<string> _addedFiles, out List<string> _removedFiles)
        {
            _addedFiles = new List<string>();
            _removedFiles = new List<string>();
            
            FRNInfo info;
            if (m_frnInfos.TryGetValue(_frn, out info) && info.IsFolder)
            {
                List<KeyValuePair<UInt64, FRNInfo>> childrenList;
                GetChildrenOf(_frn, out childrenList);

                info.ParentFRN = _newParentFrn;
                info.Name = _newName;
                ComputeFullPath(_frn, info);
                foreach (KeyValuePair<UInt64, FRNInfo> pair in childrenList)
                {
                    UInt64 childKey = pair.Key;
                    FRNInfo childInfo = pair.Value;
                    if (!childInfo.IsFolder)
                    {
                        _removedFiles.Add(childInfo.NormalizedFullPath);
                        ComputeFullPath(childKey, childInfo);
                        _addedFiles.Add(childInfo.NormalizedFullPath);
                    }
                    else
                    {
                        ComputeFullPath(childKey, childInfo);
                    }
                }
            }
        }

        /// <summary>
        /// Save current state in _parentElement
        /// </summary>
        public bool Save(XElement _parentElement)
        {
            XElement element = new XElement("FRNInfosList");
            _parentElement.Add(element);

            foreach (KeyValuePair<UInt64, FRNInfo> pair in m_frnInfos)
            {
                XElement frnElement = new XElement("FRNInfo");
                element.Add(frnElement);
                element.SetAttributeValue("RootFRN", m_folderFRN);

                frnElement.SetAttributeValue("FRN", pair.Key);
                FRNInfo info = pair.Value;
                frnElement.SetAttributeValue("ParentFRN", info.ParentFRN);
                frnElement.SetAttributeValue("Name", info.Name);
                frnElement.SetAttributeValue("FullPath", info.NormalizedFullPath);
                Debug.Assert(!string.IsNullOrEmpty(info.NormalizedFullPath));
                frnElement.SetAttributeValue("IsFolder", info.IsFolder);
            }

            return true;
        }

        /// <summary>
        /// Load current state from _parentElement
        /// </summary>
        public bool Load(XElement _parentElement)
        {
            m_frnInfos.Clear();

            XElement element = _parentElement.Element("FRNInfosList");
            if (element == null)
                return false;
            if (!UInt64.TryParse(element.Attribute("RootFRN").Value, out m_folderFRN))
                return false;

            foreach (XElement frnElement in element.Elements("FRNInfo"))
            {
                UInt64 frn;
                UInt64 parentFrn;
                string name = frnElement.Attribute("Name").Value;
                string fullPath = frnElement.Attribute("FullPath").Value;
                bool isFolder;
                if (
                    UInt64.TryParse(frnElement.Attribute("FRN").Value, out frn) &&
                    UInt64.TryParse(frnElement.Attribute("ParentFRN").Value, out parentFrn) &&
                    !string.IsNullOrEmpty(name) &&
                    !string.IsNullOrEmpty(fullPath) &&
                    Boolean.TryParse(frnElement.Attribute("IsFolder").Value, out isFolder))
                {
                    FRNInfo newInfo = new FRNInfo(name, parentFrn, isFolder);
                    newInfo.NormalizedFullPath = fullPath;
                    m_frnInfos.Add(frn, newInfo);
                }
                else
                {
                    m_frnInfos.Clear();
                    return false;
                }
            }

            return true;
        }
    }
}