// Copyright (c) 2018, Rene Lergner - @Heathcliff74xda
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

using QCEDL.NET.Logging;
using Qualcomm.EmergencyDownload.Layers.PBL.Sahara.Command;
using Qualcomm.EmergencyDownload.Transport;

namespace Qualcomm.EmergencyDownload.Layers.PBL.Sahara;

internal delegate void ReadyHandler();

public class QualcommSahara(IQualcommTransport transport)
{
    private const uint FirehoseProgrammerImageId = 13;
    private const int SaharaPacketHeaderLength = 0x08;
    private const int EndImageTxPacketLength = 0x10;
    private const int DoneResponsePacketLength = 0x0C;
    public uint DetectedDeviceSaharaVersion { get; private set; } = 2;
    public bool HasDetectedDeviceSaharaVersion { get; private set; }
    public bool IsDeviceSaharaVersionInferred { get; private set; }
    public uint? DetectedDeviceChipId { get; private set; }
    public bool IsCommandModeReady { get; private set; }
    private Dictionary<uint, string> _imageMappings = [];

    #region Packet Building Logic

    public static byte[] BuildCommandPacket(QualcommSaharaCommand saharaCommand, byte[]? commandBuffer = null)
    {
        var commandId = (uint)saharaCommand;
        var commandBufferLength = (uint)(commandBuffer?.Length ?? 0);
        var length = 0x8u + commandBufferLength;

        var packet = new byte[length];
        ByteOperations.WriteUInt32(packet, 0x00, commandId);
        ByteOperations.WriteUInt32(packet, 0x04, length);

        if (commandBuffer != null && commandBufferLength != 0)
        {
            Buffer.BlockCopy(commandBuffer, 0, packet, 0x08, commandBuffer.Length);
        }

        return packet;
    }

    private static byte[] BuildHelloResponsePacket(QualcommSaharaMode saharaMode, uint protocolVersion = 2, uint supportedVersion = 1, uint maxPacketLength = 0)
    {
        var hello = new byte[0x28];

        // Hello packet:
        // xxxxxxxx = Protocol version
        // xxxxxxxx = Supported version
        // xxxxxxxx = Max packet length
        // xxxxxxxx = Expected mode
        // 6 dwords reserved space
        ByteOperations.WriteUInt32(hello, 0x00, protocolVersion);
        ByteOperations.WriteUInt32(hello, 0x04, supportedVersion);
        ByteOperations.WriteUInt32(hello, 0x08, maxPacketLength);
        ByteOperations.WriteUInt32(hello, 0x0C, (uint)saharaMode);
        return BuildCommandPacket(QualcommSaharaCommand.HelloResponse, hello);
    }

    private static bool IsFirehoseXml(ReadOnlySpan<byte> packet)
    {
        return packet.StartsWith("<?xml"u8);
    }

    private static bool IsCompleteSaharaPacket(byte[] packet, QualcommSaharaCommand command)
    {
        return command != QualcommSaharaCommand.NoCommand &&
               Enum.IsDefined(command) &&
               packet.Length >= SaharaPacketHeaderLength &&
               ByteOperations.ReadUInt32(packet, 0x04) == packet.Length;
    }

    private static string DescribeUnexpectedHandshakePacket(
        byte[] packet,
        QualcommSaharaCommand command,
        string expectedCommand)
    {
        if (command == QualcommSaharaCommand.EndImageTx && packet.Length == EndImageTxPacketLength)
        {
            var imageId = ByteOperations.ReadUInt32(packet, 0x08);
            var status = ByteOperations.ReadUInt32(packet, 0x0C);
            return $"Expected Sahara {expectedCommand}, received {command} " +
                   $"(Image ID {imageId}, Status {DescribeSaharaStatus(status)}).";
        }

        return $"Expected Sahara {expectedCommand}, received {command}.";
    }

    private static void ThrowUnexpectedHandshakePacket(
        byte[] packet,
        QualcommSaharaCommand command,
        string expectedCommand)
    {
        var message = DescribeUnexpectedHandshakePacket(packet, command, expectedCommand);
        if (IsCompleteSaharaPacket(packet, command))
        {
            var isEndImageTx =
                command == QualcommSaharaCommand.EndImageTx &&
                packet.Length == EndImageTxPacketLength;
            throw new QualcommSaharaUnexpectedPacketException(
                message,
                command,
                [.. packet],
                isEndImageTx ? ByteOperations.ReadUInt32(packet, 0x08) : null,
                isEndImageTx ? ByteOperations.ReadUInt32(packet, 0x0C) : null);
        }

        throw new BadMessageException(message);
    }

    private static QualcommSaharaCommand ValidateAndLogPacket(byte[] packet)
    {
        LibraryLogger.Trace($"Received Sahara packet ({packet.Length} bytes): {Convert.ToHexString(packet)}");
        if (packet.Length < SaharaPacketHeaderLength)
        {
            throw new BadMessageException(
                $"Sahara packet is only {packet.Length} bytes; at least {SaharaPacketHeaderLength} bytes are required.");
        }

        var command = (QualcommSaharaCommand)ByteOperations.ReadUInt32(packet, 0x00);
        var declaredLength = ByteOperations.ReadUInt32(packet, 0x04);
        LibraryLogger.Debug(
            $"Received Sahara {command} (0x{(uint)command:X8}), packet length {packet.Length}, declared length {declaredLength}.");
        return declaredLength == packet.Length
            ? command
            : throw new BadMessageException(
                $"Sahara {command} length mismatch: received {packet.Length} bytes, header declares {declaredLength} bytes.");
    }

    private static bool HasExpectedPacketLength(byte[] packet, int expectedLength,
        QualcommSaharaCommand command)
    {
        if (packet.Length != expectedLength)
        {
            LibraryLogger.Error(
                $"Invalid Sahara {command} packet length: received {packet.Length} bytes, expected {expectedLength} bytes. Raw packet: {Convert.ToHexString(packet)}");
        }

        return packet.Length == expectedLength;
    }

    private static string DescribeSaharaStatus(uint status)
    {
        var statusCode = (QualcommSaharaStatusCode)status;
        return Enum.IsDefined(statusCode)
            ? $"{statusCode} (0x{status:X8})"
            : $"Unknown (0x{status:X8})";
    }

    private (byte[] Packet, bool HelloResponseSent) ReadInitialPacket(QualcommSaharaMode mode)
    {
        try
        {
            return (transport.GetResponse(null), false);
        }
        catch (TimeoutException) when (transport.Backend == TransportBackend.WindowsQud)
        {
            return RecoverDiscardedQudHello(mode);
        }
    }

    private (byte[] Packet, bool HelloResponseSent) RecoverDiscardedQudHello(QualcommSaharaMode mode)
    {
        LibraryLogger.Warning(
            $"Initial Sahara read timed out on Windows QUD; sending a speculative HELLO response for {mode} mode.");
        transport.SendData(BuildHelloResponsePacket(mode));
        return (transport.GetResponse(null), true);
    }

    private void RecordAdvertisedSaharaVersion(byte[] helloPacket)
    {
        DetectedDeviceSaharaVersion = ByteOperations.ReadUInt32(helloPacket, 0x08);
        HasDetectedDeviceSaharaVersion = true;
        IsDeviceSaharaVersionInferred = false;
    }

    #endregion

    #region Data Transfer Core

    private void SendData64Bit(FileStream fileStream, byte[] readDataRequest)
    {
        var offset = ByteOperations.ReadUInt64(readDataRequest, 0x10);
        var length = ByteOperations.ReadUInt64(readDataRequest, 0x18);
        var imageBuffer = new byte[length];

        if (fileStream.Position != (long)offset)
        {
            _ = fileStream.Seek((long)offset, SeekOrigin.Begin);
        }

        fileStream.ReadExactly(imageBuffer, 0, (int)length);
        transport.SendData(imageBuffer);
    }

    private void SendData(FileStream fileStream, byte[] readDataRequest)
    {
        var offset = ByteOperations.ReadUInt32(readDataRequest, 0x0C);
        var length = ByteOperations.ReadUInt32(readDataRequest, 0x10);
        var imageBuffer = new byte[length];

        if (fileStream.Position != offset)
        {
            _ = fileStream.Seek(offset, SeekOrigin.Begin);
        }

        fileStream.ReadExactly(imageBuffer, 0, (int)length);
        transport.SendData(imageBuffer);
    }

    #endregion

    #region Public Interface

    public IReadOnlyList<QualcommSaharaRamDumpRegion> CollectRamDump(
        string outputDirectory,
        string? segmentFilter = null,
        Action<QualcommSaharaRamDumpProgress>? progress = null)
    {
        return new QualcommSaharaRamDump(transport).Collect(outputDirectory, segmentFilter, progress);
    }

    public QualcommSaharaHandshakeResult ProbeCommandMode(
        byte[]? preReadHelloPacket = null,
        bool initialReadAlreadyTimedOut = false)
    {
        try
        {
            if (IsCommandModeReady)
            {
                LibraryLogger.Debug("Sahara command mode is already established; reusing the existing session.");
                return QualcommSaharaHandshakeResult.Sahara;
            }

            byte[] packet;
            var helloResponseSent = false;
            if (preReadHelloPacket != null && initialReadAlreadyTimedOut)
            {
                throw new ArgumentException(
                    "A pre-read HELLO packet cannot be combined with an already timed-out initial read.",
                    nameof(initialReadAlreadyTimedOut));
            }

            if (preReadHelloPacket != null)
            {
                LibraryLogger.Debug("Using pre-read HELLO packet for handshake.");
                packet = preReadHelloPacket;
            }
            else if (initialReadAlreadyTimedOut)
            {
                if (transport.Backend != TransportBackend.WindowsQud)
                {
                    throw new InvalidOperationException(
                        "Discarded HELLO recovery is only available for Windows QUD transports.");
                }

                (packet, helloResponseSent) = RecoverDiscardedQudHello(QualcommSaharaMode.Command);
            }
            else
            {
                LibraryLogger.Debug("Reading HELLO packet from device for handshake.");
                (packet, helloResponseSent) = ReadInitialPacket(QualcommSaharaMode.Command);
            }

            if (IsFirehoseXml(packet))
            {
                LibraryLogger.Debug("Device returned Firehose XML while probing Sahara command mode.");
                return QualcommSaharaHandshakeResult.Firehose;
            }

            if (packet.Length < sizeof(uint))
            {
                throw new BadMessageException("Initial Sahara response is shorter than a command ID.");
            }

            var command = (QualcommSaharaCommand)ByteOperations.ReadUInt32(packet, 0);
            if (helloResponseSent && command == QualcommSaharaCommand.CommandReady)
            {
                LibraryLogger.Debug("Windows QUD accepted the speculative HELLO response.");
                IsCommandModeReady = true;
                return QualcommSaharaHandshakeResult.Sahara;
            }

            if (command != QualcommSaharaCommand.Hello)
            {
                ThrowUnexpectedHandshakePacket(packet, command, "HELLO");
            }

            if (packet.Length < 0x0C)
            {
                throw new BadMessageException("Sahara HELLO packet is too short.");
            }

            RecordAdvertisedSaharaVersion(packet);
            var helloResponse = BuildHelloResponsePacket(
                QualcommSaharaMode.Command,
                DetectedDeviceSaharaVersion);
            var ready = transport.SendCommand(helloResponse, null);
            if (IsFirehoseXml(ready))
            {
                LibraryLogger.Debug("Device returned Firehose XML after the Sahara HELLO response.");
                return QualcommSaharaHandshakeResult.Firehose;
            }

            if (ready.Length < sizeof(uint))
            {
                throw new BadMessageException("Sahara command-mode response is shorter than a command ID.");
            }

            var responseId = (QualcommSaharaCommand)ByteOperations.ReadUInt32(ready, 0);
            if (responseId != QualcommSaharaCommand.CommandReady)
            {
                ThrowUnexpectedHandshakePacket(ready, responseId, "COMMAND_READY");
            }

            IsCommandModeReady = true;
            return QualcommSaharaHandshakeResult.Sahara;
        }
        catch (QualcommSaharaUnexpectedPacketException ex)
        {
            LibraryLogger.Error($"Handshake failed: {ex.Message}");
            return QualcommSaharaHandshakeResult.UnexpectedSaharaPacket;
        }
        catch (Exception ex)
        {
            LibraryLogger.Error($"Handshake failed: {ex.Message}");
            return QualcommSaharaHandshakeResult.Failed;
        }
    }

    public bool CommandHandshake(byte[]? preReadHelloPacket = null)
    {
        return ProbeCommandMode(preReadHelloPacket) == QualcommSaharaHandshakeResult.Sahara;
    }

    public void ResetSahara()
    {
        IsCommandModeReady = false;
        _ = transport.SendCommand(BuildCommandPacket(QualcommSaharaCommand.Reset), [0x08, 0x00, 0x00, 0x00]);
    }

    public void ResetStateMachine()
    {
        IsCommandModeReady = false;
        var packet = BuildCommandPacket(QualcommSaharaCommand.ResetStateMachine);
        LibraryLogger.Debug(
            $"Sending Sahara RESET_STATE_MACHINE ({packet.Length} bytes): {Convert.ToHexString(packet)}");
        transport.SendData(packet);
    }

    public void SwitchMode(QualcommSaharaMode mode)
    {
        IsCommandModeReady = false;
        var switchMode = new byte[0x04];
        ByteOperations.WriteUInt32(switchMode, 0x00, (uint)mode);
        transport.SendData(BuildCommandPacket(QualcommSaharaCommand.SwitchMode, switchMode));
    }

    #endregion

    public bool SendImage(string path)
    {
        try
        {
            if (path.EndsWith(".xml", StringComparison.OrdinalIgnoreCase))
            {
                _imageMappings = SaharaConfigParser.ParseAndValidateConfig(path);
                LibraryLogger.Debug("Multi-image mode validated.");
                LibraryLogger.Debug($"Loaded configuration, total {_imageMappings.Count} files.");
            }
            else
            {
                _imageMappings.Clear();
                _imageMappings[FirehoseProgrammerImageId] = path;
                if (!File.Exists(path))
                {
                    throw new FileNotFoundException("File not found", path);
                }
            }
        }
        catch (Exception ex)
        {
            LibraryLogger.Error($"Pre-transfer validation failed: {ex.Message}");
            return false;
        }

        var imagesTransferredCount = 0;
        uint? awaitingDoneResponseForImageId = null;
        try
        {
            var (initialPacket, helloResponseSent) = ReadInitialPacket(QualcommSaharaMode.ImageTxPending);
            if (IsFirehoseXml(initialPacket))
            {
                LibraryLogger.Debug("Device is already in Firehose mode; skipping Sahara image transfer.");
                return true;
            }

            if (initialPacket.Length < sizeof(uint))
            {
                throw new BadMessageException("Initial Sahara transfer response is shorter than a command ID.");
            }

            var pendingRequest = initialPacket;
            var initialCommand = (QualcommSaharaCommand)ByteOperations.ReadUInt32(initialPacket, 0);
            if (initialCommand == QualcommSaharaCommand.Hello)
            {
                if (initialPacket.Length < 0x0C)
                {
                    throw new BadMessageException("Sahara HELLO packet is too short.");
                }

                RecordAdvertisedSaharaVersion(initialPacket);
                transport.SendData(BuildHelloResponsePacket(QualcommSaharaMode.ImageTxPending));
                pendingRequest = null;
            }
            else if (!helloResponseSent)
            {
                LibraryLogger.Debug($"Sahara transfer started with {initialCommand} instead of HELLO.");
            }

            while (true)
            {
                byte[] request;
                try
                {
                    request = pendingRequest ?? transport.GetResponse(null);
                    pendingRequest = null;
                }
                catch (TimeoutException)
                {
                    if (imagesTransferredCount > 0)
                    {
                        LibraryLogger.Debug("Sahara timeout after successful transfers. Device likely switched to Firehose mode.");
                        return true;
                    }
                    throw;
                }

                var commandId = ValidateAndLogPacket(request);

                if (commandId == QualcommSaharaCommand.Hello)
                {
                    LibraryLogger.Debug("Device requested re-handshake (Next stage).");
                    transport.SendData(BuildHelloResponsePacket(QualcommSaharaMode.ImageTxPending));
                    continue;
                }

                if (commandId is QualcommSaharaCommand.ReadData or QualcommSaharaCommand.ReadData64Bit)
                {
                    var lastRequestedId = ByteOperations.ReadUInt32(request, 0x08);
                    if (_imageMappings.TryGetValue(lastRequestedId, out var filePath) && File.Exists(filePath))
                    {
                        LibraryLogger.Debug($"Requested ID: {lastRequestedId}, Path: {filePath}");
                        using var fs = new FileStream(filePath, FileMode.Open, FileAccess.Read, FileShare.Read);
                        if (commandId == QualcommSaharaCommand.ReadData)
                        {
                            SendData(fs, request);
                        }
                        else
                        {
                            SendData64Bit(fs, request);
                        }
                    }
                    else
                    {
                        LibraryLogger.Error($"Could not find Image ID: {lastRequestedId}");
                        return false;
                    }
                }
                else if (commandId == QualcommSaharaCommand.EndImageTx)
                {
                    if (!HasExpectedPacketLength(request, EndImageTxPacketLength, commandId))
                    {
                        return false;
                    }

                    var imageId = ByteOperations.ReadUInt32(request, 0x08);
                    var status = ByteOperations.ReadUInt32(request, 0x0C);
                    LibraryLogger.Debug(
                        $"Received END_IMAGE_TX: Image ID {imageId}, Status {DescribeSaharaStatus(status)}. Raw packet: {Convert.ToHexString(request)}");
                    if (status != (uint)QualcommSaharaStatusCode.StatusSuccess)
                    {
                        LibraryLogger.Error(
                            $"Device rejected Sahara image ID {imageId}: {DescribeSaharaStatus(status)}. DONE will not be sent.");
                        return false;
                    }

                    if (awaitingDoneResponseForImageId is uint previousImageId)
                    {
                        LibraryLogger.Error(
                            $"Received duplicate END_IMAGE_TX for image ID {imageId} while waiting for DONE_RESP for image ID {previousImageId}. DONE was already sent once; aborting to avoid an infinite protocol loop. Raw packet: {Convert.ToHexString(request)}");
                        return false;
                    }

                    LibraryLogger.Debug($"Image ID {imageId} transfer succeeded, sending DONE once.");
                    transport.SendData(BuildCommandPacket(QualcommSaharaCommand.Done));
                    awaitingDoneResponseForImageId = imageId;
                    imagesTransferredCount++;
                }
                else if (commandId == QualcommSaharaCommand.DoneResponse)
                {
                    if (!HasExpectedPacketLength(request, DoneResponsePacketLength, commandId))
                    {
                        return false;
                    }

                    var doneStatus = ByteOperations.ReadUInt32(request, 0x08);
                    var doneStatusDescription = doneStatus switch
                    {
                        0 => "Pending (0x00000000)",
                        1 => "Complete (0x00000001)",
                        _ => $"Unknown (0x{doneStatus:X8})"
                    };
                    LibraryLogger.Debug(
                        $"Received DONE_RESP for image ID {awaitingDoneResponseForImageId?.ToString(System.Globalization.CultureInfo.InvariantCulture) ?? "unknown"}: Status {doneStatusDescription}. Raw packet: {Convert.ToHexString(request)}");
                    if (awaitingDoneResponseForImageId is null)
                    {
                        LibraryLogger.Error("Received DONE_RESP without a preceding successful END_IMAGE_TX/DONE exchange.");
                        return false;
                    }

                    awaitingDoneResponseForImageId = null;
                    if (doneStatus == 1)
                    {
                        LibraryLogger.Debug(
                            $"Device reported Sahara transfer complete after {imagesTransferredCount} image(s).");
                        return true;
                    }

                    if (doneStatus != 0)
                    {
                        LibraryLogger.Error($"Device returned invalid Sahara DONE_RESP status {doneStatusDescription}.");
                        return false;
                    }

                    // Match qdl's compatibility quirk for older targets such as MSM8916.
                    if (_imageMappings.Count == 1 && _imageMappings.ContainsKey(FirehoseProgrammerImageId))
                    {
                        LibraryLogger.Warning(
                            "Device reported DONE_RESP Pending for a single image ID 13; treating it as complete for compatibility.");
                        return true;
                    }

                    LibraryLogger.Debug("Device expects additional Sahara images; waiting for the next request.");
                    continue;
                }
                else if (commandId == QualcommSaharaCommand.Done)
                {
                    LibraryLogger.Error("Unexpected Sahara DONE request received from the device; expected DONE_RESP.");
                    return false;
                }
                else
                {
                    LibraryLogger.Error($"Unexpected Sahara command received: 0x{(uint)commandId:X8}");
                    return false;
                }
            }
        }
        catch (Exception ex)
        {
            LibraryLogger.Error($"Transfer crash: {ex.Message}");
            return false;
        }
    }

    public Task<bool> LoadProgrammer(string programmerPath)
    {
        return Task.Run(() => SendImage(programmerPath));
    }

    #region Helper Methods for Device Info

    public byte[][] GetRkHs()
    {
        return Execute.GetRkHs(transport);
    }

    public byte[] GetRkh()
    {
        return Execute.GetRkh(transport);
    }

    public byte[] GetHwid()
    {
        return DetectedDeviceSaharaVersion >= 3
            ? Execute.GetV3ChipInfo(transport).ToHwidBytes()
            : Execute.GetHwid(transport);
    }

    public QualcommSaharaV3ChipInfo GetV3ChipInfo()
    {
        return DetectedDeviceSaharaVersion >= 3
            ? Execute.GetV3ChipInfo(transport)
            : throw new InvalidOperationException(
                $"Sahara CMD10 chip information requires protocol version 3 or newer; device reported version {DetectedDeviceSaharaVersion}.");
    }

    public byte[] GetSerialNumber()
    {
        var serialNumber = Execute.GetSerialNumber(transport, out var chipId);
        DetectedDeviceChipId = chipId;
        if (chipId.HasValue && !HasDetectedDeviceSaharaVersion)
        {
            DetectedDeviceSaharaVersion = 3;
            HasDetectedDeviceSaharaVersion = true;
            IsDeviceSaharaVersionInferred = true;
            LibraryLogger.Debug(
                $"Sahara HELLO version was unavailable; inferred protocol version 3 from the 8-byte SERIAL_NUM_READ response (Chip ID 0x{chipId.Value:X8}).");
        }
        else if (chipId.HasValue && DetectedDeviceSaharaVersion < 3)
        {
            LibraryLogger.Warning(
                $"Sahara protocol version {DetectedDeviceSaharaVersion} was advertised, but SERIAL_NUM_READ returned the 8-byte v3 layout (Chip ID 0x{chipId.Value:X8}). Keeping the advertised protocol version.");
        }

        return serialNumber;
    }

    #endregion
}