using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using Windows.Data.Xml.Dom;
using Windows.Foundation;
using Windows.Foundation.Metadata;
using Windows.Storage;
using FwUpdateDriverApi;
using FwUpdateDriverApi.OnBoardRetimer;

namespace FwUpdateApiSample
{
    internal struct OnBoardRetimerImage
    {
        public RetimerRoute Route { get; internal set; }
        public NvmImageDescriptor RetimerImage { get; internal set; }
    }

    internal class NvmImageDescriptor
    {
        public byte[] WholeImage;
    }
    internal struct NvmProperties
    {
        internal int StartOffset;
        internal int Size;
    }

    internal class DeviceImageContainer
    {
        const uint GRNvmPointerOffset = 0x14;

        public NvmImageDescriptor DeviceImage { get; }
        public IEnumerable<OnBoardRetimerImage> RetimerImages { get; }

        public DeviceImageContainer(string containerPath)
        {
            var container = new FileStream(containerPath, FileMode.Open, FileAccess.Read);
            ValidateContainerFormat(container);
            DeviceImage = ReadGRNvm(container); // Retrieve information about device image
            RetimerImages = ReadDBRsNvm(container); // Retrieve information about retimers
        }

        /// <summary>
        /// Reads the device NVM version from the container
        /// </summary>
        /// <param name="imageFileStream">Container content</param>
        /// <returns>Byte array representing the full NVM version</returns>
        public static byte[] GetImageFullNvmVersionFromContainer(FileStream imageFileStream)
        {
            // Const Offset
            const uint nvmStartVersionOffset = 0x9;

            var nvmVersion = new byte[2];

            // If Gr image is present in container call GetGrNvm, otherwise, call GetRetimerNvm
            var reader = DeviceContainerUtilities.IsGrPresentInContainer(imageFileStream) ? 
                         GetGrNvm(imageFileStream) : GetRetimerNvm(imageFileStream);

            reader.BaseStream.Seek(nvmStartVersionOffset, SeekOrigin.Current); // Increment position to start of nvm version
            nvmVersion = reader.ReadBytes(2);

            return nvmVersion;
        }

        /// <summary>
        /// Reads the device id from the container
        /// </summary>
        /// <param name="imageFileStream">Container content</param>
        /// <returns>ushort representing the device id</returns>
        public static ushort GetDeviceIdFromContainer(FileStream imageFileStream)
        {
            // Const Offset
            const uint deviceIdStartOffset = 0x5;

            var deviceId = new byte[2];
            // If Gr image is present in container call GetGrNvm, otherwise, call GetRetimerNvm
            var reader = DeviceContainerUtilities.IsGrPresentInContainer(imageFileStream) ?
                         GetGrNvm(imageFileStream) : GetRetimerNvm(imageFileStream);

            reader.BaseStream.Seek(deviceIdStartOffset, SeekOrigin.Current); // Increment position to start of device id
            deviceId = reader.ReadBytes(2);

            return BitConverter.ToUInt16(deviceId, 0);
        }

        /// <summary>
        /// Reads the device's vendor and model ids from the container
        /// </summary>
        /// <param name="imageFileStream">Container content</param>
        /// <returns>List of byte array representing the vendor id and model id of the device</returns>
        public static List<byte[]> GetImageDeviceInformationFromContainer(FileStream imageFileStream)
        {
            // Const Offset
            const uint nvmVendorIdOffset = 0x10;
            var ids = new List<byte[]>();

            // If Gr image is present in container call GetGrNvm, otherwise, call GetRetimerNvm
            var reader = DeviceContainerUtilities.IsGrPresentInContainer(imageFileStream) ?
                         GetGrNvm(imageFileStream) : GetRetimerNvm(imageFileStream);

            var baseRelativeAddress = reader.BaseStream.Position; // Save base relative address for further incrementation

            // Read Device Id from image
            reader.BaseStream.Seek(5, SeekOrigin.Current);
            var deviceId = BitConverter.ToUInt16(reader.ReadBytes(2), 0);

            // Read drom address from image, substracting 7 from offset because of previous read
            reader.BaseStream.Seek((DriverApiFactory.Settings).GetDromAddressNvmLocation(deviceId) - 7, SeekOrigin.Current);
            var dromAddress = BitConverter.ToUInt32(reader.ReadBytes(4), 0);

            reader.BaseStream.Position = baseRelativeAddress; // Set position of stream to initial relative address

            // Set position to the vendor id value in image
            reader.BaseStream.Seek(nvmVendorIdOffset + dromAddress, SeekOrigin.Current);

            ids.Add(reader.ReadBytes(2)); // First, read vendor id
            ids.Add(reader.ReadBytes(2)); // Second, read model id

            return ids;
        }

        private static readonly string ContainerGuid = "2B7A540B69D24F51A61DF77B4AE23ED7";

        /// <summary>
        /// Checks if given image is a container.
        /// </summary>
        /// <param name="path">Path for TBT image</param>
        /// <returns></returns>
        public static bool IsImageContainer(string path)
        {
            using (var fs = new FileStream(path, FileMode.Open, FileAccess.Read))
            {
                var guidFromImage = new byte[16];
                fs.Read(guidFromImage, 0, 16);
                var sb = new StringBuilder();
                for (var i = 0; i < guidFromImage.Length; i += 4)
                {
                    var reversedDw = guidFromImage.Skip(i).Take(4).Reverse().ToArray();
                    sb.Append(BitConverter.ToString(reversedDw).Replace("-", string.Empty));
                }

                return ContainerGuid == sb.ToString();
            }
        }

        /// <summary>
        /// Read GR nvm from container
        /// </summary>
        /// <param name="container">File stream for the container</param>
        /// <returns>NvmImageDescriptor holding GR's image</returns>
        private static NvmImageDescriptor ReadGRNvm(FileStream container)
        {
            var deviceImageDescriptor = new NvmImageDescriptor
            {
                WholeImage = ReadImage(container, GRNvmPointerOffset)    // Read the GR nvm
            };

            return deviceImageDescriptor.WholeImage == null ? null : deviceImageDescriptor;
        }

        /// <summary>
        /// Read DBRs nvm from container
        /// </summary>
        /// <param name="container">File stream for the container</param>
        /// <returns>List of OnBoardRetimerImage instances which hold the DBRs' nvm</returns>
        private static IEnumerable<OnBoardRetimerImage> ReadDBRsNvm(FileStream container)
        {
            // Variables
            var dbrDescriptorContainerSection = 0x1F;  // Location to DBR section in container
            var offsetToNextDBRNvm = 0xC;
            var onBoardRetimers = new List<OnBoardRetimerImage>();

            // Retrieve number of retimers from the container
            var retimersCount = ReadNumberOfRetimers(container);

            // Retrieve Pointer and Length values for each retimer present in container
            for (var i = 0; i < retimersCount; i++, dbrDescriptorContainerSection += offsetToNextDBRNvm)
            {
                container.Seek(dbrDescriptorContainerSection, SeekOrigin.Begin);  // Moving cursor to the DBR section's location in container
                var currentOnBoardRetimer = new OnBoardRetimerImage()
                {
                    Route = new RetimerRoute()
                    {
                        SmBusAddress = ReadSMbusAddress(container)
                    },

                    RetimerImage = new NvmImageDescriptor()
                    {
                        // Pointer to NVM is located one byte after SMBus address in container (little endian)
                        // therefore adding one to the offset
                        WholeImage = ReadImage(container, (uint)(dbrDescriptorContainerSection + 1))
                    }
                };

                onBoardRetimers.Add(currentOnBoardRetimer);
            }

            return onBoardRetimers;
        }

        /// <summary>
        /// Helper function to read an image from a container
        /// </summary>
        /// <param name="container">File stream for the container</param>
        /// <param name="nvmDetailsOffset">Offset in container for the nvm's descriptor</param>
        /// <returns>Byte array representing an image</returns>
        private static byte[] ReadImage(FileStream container, uint nvmDetailsOffset)
        {
            container.Seek(nvmDetailsOffset, SeekOrigin.Begin);
            var reader = new BinaryReader(container);

            var nvmDescriptor = DeviceContainerUtilities.ExtractNvmDescriptor(container, nvmDetailsOffset);

            // Nvm payload exists
            if (nvmDescriptor.Size > 0)
            {
                container.Seek(nvmDescriptor.StartOffset, SeekOrigin.Begin);
                return reader.ReadBytes(nvmDescriptor.Size);
            }

            return null;
        }

        /// <summary>
        /// Read SMBus address of a retimer from container
        /// </summary>
        /// <param name="container">File stream for the container</param>
        /// <returns>Byte representing SMBuss address</returns>
        private static byte ReadSMbusAddress(FileStream container)
        {
            return Convert.ToByte(container.ReadByte());
        }

        /// <summary>
        /// Read number of retimers present in the container
        /// </summary>
        /// <param name="container">File stream for the container</param>
        /// <returns>UInt16 representing the number of retimers in the container</returns>
        private static UInt16 ReadNumberOfRetimers(FileStream container)
        {
            // Variables
            const uint numOfRetimerContainerOffset = 0x10;  // Location of number of retimers value in container
            var numOfRetimers = new byte[2];

            container.Seek(numOfRetimerContainerOffset, SeekOrigin.Begin);   // Move cursor to where number of retimers value is located
            container.Read(numOfRetimers, 0, 2);
            return BitConverter.ToUInt16(numOfRetimers, 0);
        }

        /// <summary>
        /// This function validates the container's format is correct according to it's definition
        /// </summary>
        /// <param name="container">File stream for the container</param>
        private static void ValidateContainerFormat(FileStream container)
        {
            // Validate container's length
            IsContainerLengthValid(container.Length);
            
            // Extract number of retimers from container
            var numberOfRetimers = ReadNumberOfRetimers(container);

            // Extract descriptors for both GR + DBR and return offset for descriptors' section end
            var endOfDescriptorsSection = DeviceContainerUtilities.ExtractDescriptors(container, out NvmProperties GRProperties,
                                                      out Queue<NvmProperties> RetimersProperties, numberOfRetimers);

            // Validate number of retimers
            IsNumberOfRetimersValid(GRProperties.StartOffset, RetimersProperties.FirstOrDefault().StartOffset, 
                                    endOfDescriptorsSection, numberOfRetimers);

            // Validate images' sizes for both GR + DBR
            CheckImageSizes(GRProperties, RetimersProperties, container.Length);
        }


        /// <summary>
        /// This function validates the correctness of the number of retimers provided in container
        /// </summary>
        /// <param name="GRNvmStartOffset">Start offset of GR's nvm in container</param>
        /// <param name="firstDBRNvmStartOffset">Start offset of the first DBR's nvm in container</param>
        /// <param name="endOfDBRDescriptors">End of descriptors' offset in container</param>
        /// <param name="numberOfRetimers">Number of retimers provided in container</param>
        private static void IsNumberOfRetimersValid(int GRNvmStartOffset, int firstDBRNvmStartOffset, 
                                             uint endOfDBRDescriptors, uint numberOfRetimers)
        {
            // Number of retimers provided is 0 but a DBR descriptor is present in container
            // therefore, endOfDBRDescriptors' value is not located at the start of the GR nvm 
            if (numberOfRetimers == 0 && GRNvmStartOffset != 0x16)
            {
                throw new TbtException(TbtStatus.SDK_INVALID_CONTAINER_INPUT);
            }

            // If end of the DBR's descriptors in container does not equal to the start of the first DBR's nvm
            // it is a mismatch between number of retimers provided in container and actual DBRs described in container
            if (numberOfRetimers > 0 && endOfDBRDescriptors != firstDBRNvmStartOffset)
            {
                throw new TbtException(TbtStatus.SDK_NUMBER_OF_RETIMERS_MISMATCH);
            }
        }

        /// <summary>
        /// This function validates the correctness of the container's length 
        /// </summary>
        /// <param name="containerLength">Container's length</param>
        private static void IsContainerLengthValid(long containerLength)
        {
            // Container is empty except for the GUID or GUID + Version or GUID + Version + Number of Retimers
            // This function is called after IsImageContainer, therefore no need to check for length less than 16
            if (containerLength >= 16 && containerLength <= 20)
            {
                throw new TbtException(TbtStatus.SDK_INVALID_CONTAINER_INPUT);
            }
        }

        /// <summary>
        /// This function validates the sizes for GR and DBR images
        /// </summary>
        /// <param name="GRProperties">GR properties (include nvm pointer and length)</param>
        /// <param name="RetimersProperties">List of DBR properties (include nvm pointer and length for each DBR)</param>
        /// <param name="endOfContainer">Offset of the container's end</param>
        private static void CheckImageSizes(NvmProperties GRProperties, Queue<NvmProperties> RetimersProperties, long endOfContainer)
        {
            // Check GR is located last in container
            if (RetimersProperties.Any(retimer => retimer.StartOffset > GRProperties.StartOffset))
            {
                throw new TbtException(TbtStatus.SDK_INVALID_CONTAINER_FORMAT);
            }

            // Check if actual size of the GR's nvm is equal to the size provided - should be equal to the container length's value
            if (GRProperties.StartOffset + GRProperties.Size != endOfContainer)
            {
                throw new TbtException(TbtStatus.SDK_INVALID_CONTAINER_INPUT);
            }

            // Check actual DBR nvm size is equal to the size provided in the container for every DBR
            while (RetimersProperties.Any())
            {
                var retimer = RetimersProperties.Dequeue();
                var calculatedDBRNvmSize = retimer.StartOffset + retimer.Size;

                if ((GRProperties.StartOffset < calculatedDBRNvmSize) ||
                    RetimersProperties.Any(r => r.StartOffset < calculatedDBRNvmSize))
                {
                    throw new TbtException(TbtStatus.SDK_INVALID_CONTAINER_INPUT);
                }
            }
        }

        private static BinaryReader GetGrNvm(FileStream imageFileStream)
        {
            var grDescriptor = DeviceContainerUtilities.ExtractNvmDescriptor(imageFileStream, GRNvmPointerOffset);
            return DeviceContainerUtilities.GetRelativeNvmFromContainer(imageFileStream, grDescriptor.StartOffset);
        }

        private static BinaryReader GetRetimerNvm(FileStream imageFileStream)
        {
            const uint dbrDescriptorsStartOffset = 0x20;    // Location to the start of dbr's descriptors in container
            var grDescriptor = DeviceContainerUtilities.ExtractNvmDescriptor(imageFileStream, dbrDescriptorsStartOffset);
            return DeviceContainerUtilities.GetRelativeNvmFromContainer(imageFileStream, grDescriptor.StartOffset);

        }
    }

    internal static class DeviceContainerUtilities
    {
        private static readonly uint _grNvmPointerOffset = 0x14;    // Location of gr's nvm pointer in container
        
        /// <summary>
        /// Read nvm descriptor from container (includes pointer to the nvm in container and nvm's size)
        /// </summary>
        /// <param name="container">File stream for the container</param>
        /// <param name="offsetNvmProperties">Offset in container for the nvm's descriptor</param>
        /// <returns>NvmProperties struct holding the pointer to the nvm and it's size</returns>
        internal static NvmProperties ExtractNvmDescriptor(FileStream container, uint offsetNvmProperties)
        {
            var reader = new BinaryReader(container);
            reader.BaseStream.Seek(offsetNvmProperties, SeekOrigin.Begin);
            return new NvmProperties()
            {
                StartOffset = reader.ReadInt32(),
                Size = reader.ReadInt32()
            };
        }

        /// <summary>
        /// This function extracts nvm descriptors for both GR and DBR
        /// </summary>
        /// <param name="container">File stream for the container</param>
        /// <param name="GRProperties">An out NvmProperties instance representing GR</param>
        /// <param name="RetimersProperties">An out NvmProperties instance representing list of DBR</param>
        /// <param name="numOfRetimers">Number of retimers provided in container</param>
        /// <returns>Offset for the end of descriptors section in container</returns>
        internal static uint ExtractDescriptors(FileStream container, out NvmProperties GRProperties,
                                        out Queue<NvmProperties> RetimersProperties, uint numOfRetimers)
        {
            // Variables
            uint currDBRDescriptorOffset = 0x20;  // Location to start of the DBR section in container
            RetimersProperties = new Queue<NvmProperties>();    // Use Queue in order to keep track of the order of which the descriptors appeared in container

            GRProperties = ExtractNvmDescriptor(container, _grNvmPointerOffset);

            // Extract Pointer and Length for every DBR described in container
            for (var i = 0; i < numOfRetimers; i++, currDBRDescriptorOffset += 8)
            {
                var currentDBRDescriptor = ExtractNvmDescriptor(container, currDBRDescriptorOffset);
                RetimersProperties.Enqueue(currentDBRDescriptor);
            }

            return currDBRDescriptorOffset;
        }

        internal static BinaryReader GetRelativeNvmFromContainer(FileStream imageFileStream, int pointerPosition)
        {
            var reader = new BinaryReader(imageFileStream);

            // Set stream's position to the start of the nvm
            reader.BaseStream.Seek(pointerPosition, SeekOrigin.Begin);

            // Read header's size and moving cursor to the location of the relative section in nvm
            var headerLength = reader.ReadInt32() & Utilities.RelativeNvmOffsetRetrieveMask;

            // Position got advanced by 4 bytes after ReadInt32(), therefore reducing 4 bytes in calculation
            reader.BaseStream.Seek(headerLength - 4, SeekOrigin.Current);

            return reader;
        }

        internal static bool IsGrPresentInContainer(FileStream imageFileStream)
        {
            var grDescriptor = ExtractNvmDescriptor(imageFileStream, _grNvmPointerOffset);
            return grDescriptor.Size != 0;  // If size of gr in container is not zero - gr is present in container
        }
    }
}