/*******************************************************************************
* Copyright (C) 2014 - 2017 Intel Corp. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
*  - Redistributions of source code must retain the above copyright notice,
*    this list of conditions and the following disclaimer.
*
*  - Redistributions in binary form must reproduce the above copyright notice,
*    this list of conditions and the following disclaimer in the documentation
*    and/or other materials provided with the distribution.
*
*  - Neither the name of Intel Corp. nor the names of its
*    contributors may be used to endorse or promote products derived from this
*    software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS''
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
* ARE DISCLAIMED. IN NO EVENT SHALL Intel Corp. OR THE CONTRIBUTORS
* BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
* CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
* SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
* INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
* CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
* ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
* POSSIBILITY OF SUCH DAMAGE.
*******************************************************************************/

using System;
using System.Collections.Generic;
using System.Globalization;
using System.Linq;
using System.IO;
using FwUpdateDriverApi;
using System.Text;
using System.Threading;

namespace FwUpdateApiSample
{
    /// <summary>
    /// This class is the interface to a device controller
    ///
    /// This, with its base class, is the main interface of this API module for applications
    /// to interface with device controllers.
    /// The main role of this class is to wrap the interface to driver but it includes also the
    /// interface for validating image compatibility with the current controller.
    /// </summary>
    public class SdkTbtDevice : SdkTbtBase
    {
        private IDriverDevice _deviceIf;
        private readonly DeviceParams _deviceParams;

        /// <summary>
        /// C-tor
        /// </summary>
        public SdkTbtDevice(IDriverDevice deviceIf) : base(deviceIf)
        {
            _deviceIf = deviceIf;
            _deviceParams = deviceIf.getParams();
        }

        public override bool HasSharedNvmWithPd => 
            Utilities.IsLegacyDevice(ushort.Parse(
                _deviceParams.DeviceId.Replace("0x", string.Empty), 
                NumberStyles.HexNumber));

        /// <summary>
        /// This property is used to differentiate between the power down messages
        /// given to user for host controller or device controller after update in
        /// cases where the PD version was updated. 
        /// </summary>
        /// <returns>Power down necessary for device controller message string</returns>
        public override string GetNeedPowerDownMessage
        {
            get { return Resources.DeviceNeedsPowerDownMessage; }
        }

        /// <summary>
        /// Device controller UUID property as returned by the driver
        /// </summary>
        public string UUID
        {
            get { return _deviceParams.UUID; }
        }

        /// <summary>
        /// ID of the host controller the device is connected to as returned by the driver
        /// </summary>
        public string ControllerId
        {
            get { return _deviceParams.ControllerId; }
        }

        /// <summary>
        /// 0-based index of the port in the host controller the device is connected to as
        /// returned by the driver
        /// </summary>
        public UInt32 PortNum
        {
            get { return _deviceParams.PortNum; }
        }

        /// <summary>
        /// 1-based index of the device position in the port the device is connected to as
        /// returned by the driver
        /// </summary>
        public UInt32 PositionInChain
        {
            get { return _deviceParams.PositionInChain; }
        }

        /// <summary>
        /// Device vendor name as returned by the driver
        /// </summary>
        public string VendorName
        {
            get { return _deviceParams.VendorName; }
        }

        /// <summary>
        /// Device model name as returned by the driver
        /// </summary>
        public string ModelName
        {
            get { return _deviceParams.ModelName; }
        }

        /// <summary>
        /// Device vendor ID as returned by the driver
        /// </summary>
        public UInt16 VendorId
        {
            get { return _deviceParams.VendorId; }
        }

        /// <summary>
        /// Device model ID as returned by the driver
        /// </summary>
        public UInt16 ModelId
        {
            get { return _deviceParams.ModelId; }
        }

        /// <summary>
        /// Controller enumeration out of n controllers in the device (as marked in the controller's NVM)
        /// </summary>
        public byte ControllerNum
        {
            get { return _deviceParams.ControllerNum; }
        }

        /// <summary>
        /// Total amount of controllers in the device
        /// </summary>
        public byte NumOfControllers
        {
            get { return _deviceParams.NumOfControllers; }
        }

        /// <summary>
        /// True if the device is updatable according to the information from the driver
        /// </summary>
        /// <remarks>
        /// If the controller doesn't support device FW update, the device appears like not
        /// updatable even if it supports FW update.
        /// </remarks>
        public bool Updatable
        {
            get { return _deviceParams.Updatable; }
        }

        /// <summary>
        /// CIO link speed (in Gbps) of the device
        /// </summary>
        public byte LinkSpeed
        {
            get { return _deviceParams.LinkSpeed; }
        }

        /// <summary>
        /// True if device is connected through Usb, false otherwise
        /// </summary>
        public bool ConnectedThroughUsb
        {
            get { return _deviceParams.ConnectedThroughUsb; }
        }

        /// <summary>
        /// True if the device is a usb4 device or a legacy device
        /// </summary>
        public bool IsUsb4Device
        {
            get { return _deviceParams.IsUsb4Device; }
        }
        /// <summary>
        /// True if the device has read drom support, false otherwise
        /// </summary>
        public bool HasReadDromSupport
        {
            get { return _deviceParams.HasReadDromSupport; }
        }

        /// <summary>
        /// True if device is connected through SW CM, false otherwise
        /// </summary>
        public bool ConnectedThroughSwCm
        {
            get { return _deviceParams.ConnectedThroughSwCm; }
        }

        /// <summary>
        /// Returns if the device has a known topology. If it doesn't it shall be displayed flat.
        /// </summary>
        public bool IsTopologyKnown => !(ConnectedThroughSwCm || ConnectedThroughUsb);

        /// <summary>
        /// Returns all the connected devices.
        /// </summary>
        /// <returns></returns>
        public static Dictionary<string, SdkTbtDevice> GetDevices()
        {
            Dictionary<string, IDriverDevice> devicesInterfaces = null;
            Utilities.WrapWithConvertToTbtException(() => devicesInterfaces = DriverApiFactory.GetDevices(Logger.Instance));
            var devices = new Dictionary<string, SdkTbtDevice>();
            foreach (var device in devicesInterfaces)
            {
                var tbtDevice = new SdkTbtDevice(device.Value);
                Logger.Instance.LogInfo(tbtDevice.ToString());
                devices.Add(device.Key, tbtDevice);
            }
            return devices;
        }


        /// <summary>
        /// Validates FW image file compatibility with the current controller
        /// *** IMPORTANT NOTE ***
        /// The validation works only on Intel images and devices.
        /// Other USB4 vendors should implement the validation functionality using the SDK provided APIs.
        /// </summary>
        /// <param name="path">Path to FW image file to check</param>
        /// <exception>Throws exception in case of incompatibility error or other errors</exception>
        public override void ValidateImage(string path)
        {
            if (!File.Exists(path))
            {
                throw new TbtException(TbtStatus.SDK_FILE_NOT_FOUND);
            }
            if (!Updatable)
            {
                throw new TbtException(TbtStatus.SDK_DEVICE_NOT_SUPPORTED);
            }

            var deviceId = Convert.ToUInt16(_deviceParams.DeviceId, 16);

            if (!Utilities.IsIntelDevice(deviceId))
            {
                throw new TbtException(TbtStatus.DEVICE_NOT_SUPPORTED);
            }

            if (DeviceImageContainer.IsImageContainer(path))
            {
                ValidateContainer(path);
            }
            else 
            {
                byte[] imageBuffer;
                using (var fs = new FileStream(path, FileMode.Open, FileAccess.Read))
                {
                    var reader = new BinaryReader(fs);
                    imageBuffer = reader.ReadBytes((int)fs.Length);
                }

                if (!RetimerImageUtilities.IsDbrImage(path))
                {
                    validateSecurityRevision(imageBuffer);
                }

                if (_deviceIf.getParams().IsUsb4Device)
                {
                    // Validate image in case it is a Dbr image
                    if (RetimerImageUtilities.IsDbrImage(path))
                    {
                        var retimerAddress = RetimerImageUtilities.ExtractRetimerAddress(path); // Extract SmBus address from image
                        var dbrNvmDescriptor = new OnBoardRetimerImage
                        {
                            RetimerImage = new NvmImageDescriptor { WholeImage = imageBuffer },
                            Route = new RetimerRoute { SmBusAddress = retimerAddress }
                        };

                        if (!IntelUsb4Validator.DoesMatch(dbrNvmDescriptor.RetimerImage, _deviceIf, dbrNvmDescriptor.Route))
                        {
                            throw new TbtException(TbtStatus.SDK_IMAGE_FOR_RETIMER_MISMATCH);
                        }

                        return;
                    }

                    var nvmDescriptor = new NvmImageDescriptor {WholeImage = imageBuffer};
                    if(!IntelUsb4Validator.DoesMatch(nvmDescriptor, _deviceIf))
                    {
                        throw new TbtException(TbtStatus.SDK_IMAGE_FOR_DEVICE_MISMATCH);
                    }
                }
                else
                {
                    _ValidateTbt3Image(imageBuffer);
                } 
            }
        }

        private void _ValidateTbt3Image(byte[] imageBuffer)
        {
            var deviceFwInfo = new DeviceFwInfo(new ControllerFwInfoSource(this));
            var fileFwInfo = new DeviceFwInfo(new FileFwInfoSource(imageBuffer));

            if (deviceFwInfo.Info.Generation != fileFwInfo.Info.Generation)
            {
                throw new TbtException(TbtStatus.SDK_HW_GENERATION_MISMATCH);
            }
            if (deviceFwInfo.Info.Type != fileFwInfo.Info.Type)
            {
                throw new TbtException(TbtStatus.SDK_PORT_COUNT_MISMATCH);
            }

            var validator = new DeviceImageValidator(this,
                                                     imageBuffer,
                                                     deviceFwInfo.GetSectionInfo(),
                                                     fileFwInfo.GetSectionInfo(),
                                                     deviceFwInfo.Info);
            validator.Validate();
        }

        /// <summary>
        /// Validates FW and image file security level.
        /// <param name="imageBuffer">byte array of the image file to check</param>
        /// <exception>Throws exception in case of mismatch error or other errors</exception>
        private void validateSecurityRevision(byte[] imageBuffer)
        {
            var deviceFwInfo = new DeviceFwInfo(new ControllerFwInfoSource(this));
            var fileFwInfo = new DeviceFwInfo(new FileFwInfoSource(imageBuffer));

            SecurityRevisionImageValidator securityRevisionImageValidator = new SecurityRevisionImageValidator(this,
                                                                                                               imageBuffer,
                                                                                                               deviceFwInfo.GetBaseSectionInfo(),
                                                                                                               fileFwInfo.GetBaseSectionInfo(),
                                                                                                               deviceFwInfo.Info);

            securityRevisionImageValidator.Validate();
        }

        /// <summary>
        /// The function to validate the container using USB4 validation on DROM.
        /// </summary>
        /// <param name="containerImagePath"></param>
        public void ValidateContainer(string containerImagePath)
        {
            var container = new DeviceImageContainer(containerImagePath);
            if (container.DeviceImage != null)  // there is a device image in the container
            {
                Logger.Instance.LogInfo($"Found device image in the image container, validating");
                if (!IntelUsb4Validator.DoesMatch(container.DeviceImage, _deviceIf))
                {
                    throw new TbtException(TbtStatus.SDK_IMAGE_FOR_DEVICE_MISMATCH);
                }
            }
            if (container.RetimerImages.Count() > _deviceIf.OnBoardRetimers.Count())
            {
                throw new TbtException(TbtStatus.SDK_NUMBER_OF_RETIMERS_MISMATCH);
            }
            foreach (var retimerInContainer in container.RetimerImages)
            {
                if (!IntelUsb4Validator.DoesMatch(retimerInContainer.RetimerImage, _deviceIf, retimerInContainer.Route))
                {
                    throw new TbtException(TbtStatus.SDK_IMAGE_FOR_RETIMER_MISMATCH);
                }
            }

            // All the validation on the container is done.
            Logger.Instance.LogInfo($"Validation on device's {_deviceParams.ModelName} container is done");
        }

        public override void UpdateFirmware(UInt32 bufferSize, byte[] buffer)
        {
            Logger.Instance.LogInfo($"Start update {ModelName}'s firmware, buffer size is {bufferSize}");
            Utilities.WrapWithConvertToTbtException(() => _deviceIf.UpdateFirmware(bufferSize, buffer));
        }

        public override void UpdateFirmwareFromFile(string filename)
        {
            Logger.Instance.LogInfo($"Update firmware from file with file name - {filename} for device {_deviceParams.ModelName}");

            if (Utilities.HasLinks(filename))
            {
                throw new TbtException(TbtStatus.SDK_GENERAL_ERROR_CODE, $"The provided file - {filename} contains links. Therefore denied for update");
            }

            if (DeviceImageContainer.IsImageContainer(filename))
            {
                UpdateFirmwareFromContainer(filename);
            }
            else
            {
                if (RetimerImageUtilities.IsDbrImage(filename))
                {
                    UpdateFirmwareForRetimer(filename);
                }
                else
                {
                    // this part is to show how to use the two part NVM update flow.
                    if (_deviceIf.IsTwoStageFirmwareUpdateSupported)
                    {
                        using (var fs = new FileStream(filename, FileMode.Open, FileAccess.Read))
                        {
                            var reader = new BinaryReader(fs);
                            var imageBuffer = reader.ReadBytes((int)fs.Length);
                            Logger.Instance.LogInfo("First stage is to write FW");
                            _deviceIf.WriteFirmware(imageBuffer);
                            // Allow some time for controller to get to RTD3.
                            // This simulates the flow, with delayed authentication phase.
                            // Thread.Sleep(TimeSpan.FromSeconds(30));
                            Logger.Instance.LogInfo("Second stage is to authenticate it");
                            _deviceIf.AuthenticateFirmware();
                        }
                    }
                    else
                    {
                        base.UpdateFirmwareFromFile(filename);
                    }

                }
            }
        }

        public override string GetFwUpdateMethodString()
        {
            if (ConnectedThroughUsb)
            {
                return "Usb";
            }
            else if (ConnectedThroughSwCm)
            {
                return "Sw-Cm";
            }

            return "Tbt";
        }

        public void UpdateFirmwareFromContainer(string filename)
        {
            var container = new DeviceImageContainer(filename);
            var numberOfImagesToUpdate = (container.DeviceImage != null ? 1 : 0) + container.RetimerImages.Count();
            FwUpdateProgressFactory.GetProgressWatcher().NumOfImagesToUpdate = (uint)numberOfImagesToUpdate;

            // Start updating retimers if exist
            foreach (var retimerInContainer in container.RetimerImages)
            {
                var retimer = _deviceIf.OnBoardRetimers.FirstOrDefault(r => r.Route.IsSame(retimerInContainer.Route));
                if (retimer == null)
                {
                    throw new TbtException(TbtStatus.SDK_INVALID_CONTAINER_INPUT);
                }
                Logger.Instance.LogInfo($"Found retimer with address - {retimer.Route} image in the image container, starting update");
                Utilities.WrapWithConvertToTbtException(() => retimer.UpdateFirmware((uint)retimerInContainer.RetimerImage.WholeImage.Length, retimerInContainer.RetimerImage.WholeImage));
            }

            if (container.DeviceImage != null)  // there is a device image in the container
            {
                Logger.Instance.LogInfo("Found device image in the image container, starting update");
                Utilities.WrapWithConvertToTbtException(() => _deviceIf.UpdateFirmware((uint)container.DeviceImage.WholeImage.Length, container.DeviceImage.WholeImage));
            }

            Logger.Instance.LogInfo($"Device - {ModelName}: Finished firmware update for container");
        }
        
        /// <summary>
        /// Update firmware of a retimer
        /// </summary>
        /// <param name="filename">Image FW path</param>
        public void UpdateFirmwareForRetimer(string filename)
        {
            // Extract the address of the retimer to be updated
            var retimerAddress = RetimerImageUtilities.ExtractRetimerAddress(filename);
            // Find the requested retimer from the current Gr's retimers list
            var retimer = _deviceIf.OnBoardRetimers.FirstOrDefault(r => r.Route.SmBusAddress == retimerAddress);

            if (retimer == null)
            {
                throw new TbtException(TbtStatus.SDK_RETIMER_NOT_FOUND, $"Retimer with given route - {retimerAddress} not present in the device");
            }

            using (var fs = new FileStream(filename, FileMode.Open, FileAccess.Read))
            {
                var reader = new BinaryReader(fs);
                var imageBuffer = reader.ReadBytes((int)fs.Length);
                Logger.Instance.LogInfo($"Found retimer with address - {retimer.Route} image in the image container, starting update");
                Utilities.WrapWithConvertToTbtException(() => retimer.UpdateFirmware((UInt32)imageBuffer.Length, imageBuffer));
            }
        }

        /// <summary>
        /// Reads the current full NVM version (major.minor) from the retimer
        /// </summary>
        /// <param name="smBusAddress">SmBus address</param>
        /// <returns>NVM version that was read from the firmware (in major.minor format)</returns>
        public string GetRetimerCurrentFullNvmVersion(byte smBusAddress)
        {
            // Read the Nvm version of the retimer which matches the smbus adress provided
            return _deviceIf.OnBoardRetimers.FirstOrDefault(retimer => retimer.Route.SmBusAddress == smBusAddress).NvmVersion.ToString();
        }
        
        public override string ToString()
        {
            var sb = new StringBuilder();
            sb.Append('-', 80).Append("\n");
            sb.AppendLine($"Device with UUID: {UUID}");
            sb.AppendLine($"Controller ID: {ControllerId ?? "NA"}");
            sb.AppendLine($"PortNum: {PortNum}");
            sb.AppendLine($"PositionInChain: {PositionInChain}");
            sb.AppendLine($"VendorName: {VendorName}");
            sb.AppendLine($"ModelName: {ModelName}");
            sb.AppendLine($"ModelId: {ModelId}");
            sb.AppendLine($"VendorId: {VendorId}");
            sb.AppendLine($"ControllerNum: {ControllerNum}");
            sb.AppendLine($"NumOfControllers: {NumOfControllers}");
            sb.AppendLine($"Updatable: {Updatable}");
            sb.AppendLine($"LinkSpeed: {LinkSpeed}");
            sb.AppendLine($"Number of retimers onboard - {_deviceIf.OnBoardRetimers.Count()}");
            return sb.ToString();
        }
    }
}