using System.Buffers;
using System.Net;
using System.Net.Sockets;
using System.Runtime.InteropServices;
using System.Diagnostics;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using Fleet.Core;

namespace MoShiro.Measurement;

internal delegate void MsQuicDatagramCallback(MsQuicConnection connection, ReadOnlySpan<byte> datagram);

/// <summary>A single MsQuic connection with connection-bound datagram delivery.</summary>
internal sealed class MsQuicConnection : IAsyncDisposable, IDatagramConnection
{
    private const uint DatagramReceiveEnabled = 0x0500000D;
    private const uint ClientCredentialFlags = 0x00000001 | 0x00000004 | 0x00000010 | 0x00004000;
    private const uint CertificateCredentialNone = 0;
    private const uint CertificateCredentialContext = 3;
    private const uint CertificateCredentialPkcs12 = 6;
    private const uint ConnectionEventConnected = 0;
    private const uint ConnectionEventShutdownByTransport = 1;
    private const uint ConnectionEventShutdownByPeer = 2;
    private const uint ConnectionEventShutdownComplete = 3;
    private const uint ConnectionEventPeerAddressChanged = 5;
    private const uint ConnectionEventDatagramStateChanged = 10;
    private const uint ConnectionEventDatagramReceived = 11;
    private const uint ConnectionEventDatagramSendStateChanged = 12;
    private const uint ConnectionEventPeerCertificateReceived = 15;
    private const int DatagramStateFinal = 3;
    private const int QuicAddressBytes = 28;
    private static readonly MsQuicApi.ConnectionCallback s_callback = OnEvent;

    private readonly object _sendGate = new();
    private readonly TaskCompletionSource _connected = new(TaskCreationOptions.RunContinuationsAsynchronously);
    private readonly TaskCompletionSource _shutdown = new(TaskCreationOptions.RunContinuationsAsynchronously);
    private readonly GCHandle _self;
    private readonly byte[]? _expectedPeerSpki;
    private readonly Action<MsQuicConnection>? _closed;
    private readonly Action<IPEndPoint>? _peerAddressChanged;
    private readonly INodeLog _log;
    private readonly Dictionary<nint, SendState> _sends = [];
    private nint _registration;
    private nint _configuration;
    private nint _handle;
    private int _maxDatagramBytes;
    private int _sendEnabled;
    private int _disposeStarted;
    private double? _handshakeDurationMs;

    private MsQuicConnection(
        nint handle,
        nint registration,
        nint configuration,
        byte[]? expectedPeerSpki,
        MsQuicDatagramCallback? datagramHandler,
        INodeLog log,
        Action<MsQuicConnection>? closed,
        Action<IPEndPoint>? peerAddressChanged)
    {
        _handle = handle;
        _registration = registration;
        _configuration = configuration;
        _expectedPeerSpki = expectedPeerSpki;
        DatagramHandler = datagramHandler;
        _log = log;
        _closed = closed;
        _peerAddressChanged = peerAddressChanged;
        _self = GCHandle.Alloc(this);
    }

    public MsQuicDatagramCallback? DatagramHandler { get; set; }

    public IPEndPoint? RemoteEndPoint { get; private set; }

    public IPEndPoint? LocalEndPoint => Volatile.Read(ref _handle) is nint handle && handle != 0
        ? ReadEndpoint(MsQuicApi.GetAddress(handle, 0x05000001))
        : null;

    public nint Handle => Volatile.Read(ref _handle);

    public IPEndPoint? RefreshRemoteEndPoint()
    {
        if (RemoteEndPoint is not null || Handle == 0)
        {
            return RemoteEndPoint;
        }

        try
        {
            RemoteEndPoint = ReadEndpoint(MsQuicApi.GetAddress(Handle, 0x05000002));
        }
        catch (InvalidOperationException)
        {
            // Address isn't available until the handshake finishes.
        }

        return RemoteEndPoint;
    }

    public bool IsReady => _connected.Task.IsCompletedSuccessfully && Volatile.Read(ref _sendEnabled) != 0;

    public int MaxDatagramBytes => Volatile.Read(ref _maxDatagramBytes);

    public double? HandshakeDurationMs => _handshakeDurationMs;

    public bool TryGetStatistics(out MsQuicApi.QuicStatisticsV2 statistics)
    {
        nint handle = Volatile.Read(ref _handle);
        if (handle == 0)
        {
            statistics = default;
            return false;
        }

        try
        {
            statistics = MsQuicApi.GetStatisticsV2(handle);
            return true;
        }
        catch (InvalidOperationException)
        {
            statistics = default;
            return false;
        }
    }

    public static async Task<MsQuicConnection> ConnectAsync(
        IPAddress address,
        int port,
        int localPort,
        byte[] expectedPeerSpki,
        MsQuicDatagramCallback datagramHandler,
        INodeLog log,
        CancellationToken ct)
    {
        nint registration = 0;
        nint configuration = 0;
        MsQuicConnection? connection = null;
        long handshakeStarted = 0;
        try
        {
            log.Info("opening MsQuic client registration");
            registration = MsQuicApi.RegistrationOpen("moshiro-probe");
            log.Info("opening MsQuic client configuration");
            configuration = MsQuicApi.ConfigurationOpen(registration, QuicAlpn.Value);
            log.Info("loading MsQuic client credentials");
            LoadClientCredentials(configuration);
            var endpoint = new IPEndPoint(address.AddressFamily == AddressFamily.InterNetworkV6 ? IPAddress.IPv6Any : IPAddress.Any, localPort);
            SocketAddress serialized = endpoint.Serialize();
            byte[] localAddress = new byte[QuicAddressBytes];
            for (int i = 0; i < serialized.Size; i++)
            {
                localAddress[i] = serialized[i];
            }
            ushort nativeFamily = ReadAddressFamily(localAddress);
            nint name = Marshal.StringToCoTaskMemUTF8(address.ToString());
            try
            {
                nint handle = MsQuicApi.ConnectionOpen(registration, s_callback, 0);
                log.Info("opened MsQuic client connection");
                connection = new MsQuicConnection(handle, registration, configuration, expectedPeerSpki, datagramHandler, log, null, null);
                registration = 0;
                configuration = 0;
                MsQuicApi.SetContext(handle, GCHandle.ToIntPtr(connection._self));
                log.Info("enabling MsQuic datagram receive");
                MsQuicApi.SetParam(handle, DatagramReceiveEnabled, 1);
                log.Info("binding MsQuic local address", ("local_port", localPort));
                MsQuicApi.ConnectionBind(handle, localAddress);
                log.Info("starting MsQuic QUIC handshake", ("remote", address), ("port", port), ("local_port", localPort));
                handshakeStarted = Stopwatch.GetTimestamp();
                MsQuicApi.ConnectionStart(handle, connection._configuration, nativeFamily, name, checked((ushort)port));
                log.Info("MsQuic connection start returned");
            }
            finally
            {
                Marshal.FreeCoTaskMem(name);
            }

            await connection._connected.Task.WaitAsync(TimeSpan.FromSeconds(10), ct);
            connection._handshakeDurationMs = Stopwatch.GetElapsedTime(handshakeStarted).TotalMilliseconds;
            return connection;
        }
        catch
        {
            if (connection is not null)
            {
                await connection.DisposeAsync();
            }

            if (configuration != 0)
            {
                MsQuicApi.ConfigurationClose(configuration);
            }

            if (registration != 0)
            {
                MsQuicApi.RegistrationClose(registration);
            }

            throw;
        }
    }

    public static MsQuicConnection Accept(
        nint handle,
        nint configuration,
        MsQuicDatagramCallback datagramHandler,
        INodeLog log,
        Action<MsQuicConnection> closed,
        Action<IPEndPoint>? peerAddressChanged)
    {
        var connection = new MsQuicConnection(handle, 0, 0, null, datagramHandler, log, closed, peerAddressChanged);
        MsQuicApi.SetContext(handle, GCHandle.ToIntPtr(connection._self));
        MsQuicApi.SetCallbackHandler(handle, s_callback, GCHandle.ToIntPtr(connection._self));
        MsQuicApi.ConnectionSetConfiguration(handle, configuration);
        MsQuicApi.SetParam(handle, DatagramReceiveEnabled, 1);
        return connection;
    }

    public bool Send(ReadOnlySpan<byte> datagram)
    {
        nint handle = Volatile.Read(ref _handle);
        if (handle == 0 || !IsReady || datagram.Length > MaxDatagramBytes)
        {
            return false;
        }

        byte[] buffer = ArrayPool<byte>.Shared.Rent(datagram.Length);
        datagram.CopyTo(buffer);
        GCHandle pinned = GCHandle.Alloc(buffer, GCHandleType.Pinned);
        nint nativeBuffer = Marshal.AllocHGlobal(Marshal.SizeOf<MsQuicApi.QuicBuffer>());
        var state = new SendState(buffer, pinned, nativeBuffer);
        GCHandle stateHandle = GCHandle.Alloc(state);
        nint context = GCHandle.ToIntPtr(stateHandle);
        lock (_sendGate)
        {
            if (_handle == 0)
            {
                stateHandle.Free();
                state.Release();
                return false;
            }

            _sends.Add(context, state);
        }

        var nativeBufferValue = new MsQuicApi.QuicBuffer
        {
            Length = (uint)datagram.Length,
            Buffer = pinned.AddrOfPinnedObject(),
        };
        Marshal.StructureToPtr(nativeBufferValue, nativeBuffer, false);
        uint status = MsQuicApi.DatagramSend(handle, nativeBuffer, context);
        if (MsQuicApi.Succeeded(status))
        {
            return true;
        }

        ReleaseSend(context);
        return false;
    }

    public async ValueTask DisposeAsync()
    {
        if (Interlocked.Exchange(ref _disposeStarted, 1) == 0 && Volatile.Read(ref _handle) is nint handle && handle != 0)
        {
            MsQuicApi.ConnectionShutdown(handle);
        }

        if (Volatile.Read(ref _handle) != 0)
        {
            await _shutdown.Task.ConfigureAwait(false);
        }
    }

    public Task WaitForShutdownAsync() => _shutdown.Task;

    private static void LoadClientCredentials(nint configuration)
    {
        var credential = new MsQuicApi.CredentialConfig { Type = CertificateCredentialNone, Flags = ClientCredentialFlags };
        if (OperatingSystem.IsWindows())
        {
            credential.Flags &= ~0x00004000u;
        }

        MsQuicApi.ConfigurationLoadCredential(configuration, ref credential);
    }

    internal static void LoadServerCredentials(nint configuration, X509Certificate2 certificate)
    {
        var credential = new MsQuicApi.CredentialConfig();
        GCHandle pfxPin = default;
        nint pkcs12 = 0;
        try
        {
            if (OperatingSystem.IsWindows())
            {
                credential.Type = CertificateCredentialContext;
                credential.Certificate = certificate.Handle;
            }
            else
            {
                byte[] pfxBytes = certificate.Export(X509ContentType.Pkcs12);
                pfxPin = GCHandle.Alloc(pfxBytes, GCHandleType.Pinned);
                var pfxCertificate = new MsQuicApi.Pkcs12Certificate
                {
                    Asn1Blob = pfxPin.AddrOfPinnedObject(),
                    Asn1BlobLength = (uint)pfxBytes.Length,
                };
                pkcs12 = Marshal.AllocHGlobal(Marshal.SizeOf<MsQuicApi.Pkcs12Certificate>());
                Marshal.StructureToPtr(pfxCertificate, pkcs12, false);
                credential.Type = CertificateCredentialPkcs12;
                credential.Flags = 0x00004000;
                credential.Certificate = pkcs12;
            }

            MsQuicApi.ConfigurationLoadCredential(configuration, ref credential);
        }
        finally
        {
            if (pkcs12 != 0)
            {
                Marshal.FreeHGlobal(pkcs12);
            }

            if (pfxPin.IsAllocated)
            {
                pfxPin.Free();
            }
        }
    }

    private static uint OnEvent(nint handle, nint context, nint eventPointer)
    {
        try
        {
            var connection = (MsQuicConnection)GCHandle.FromIntPtr(context).Target!;
            uint eventType = unchecked((uint)Marshal.ReadInt32(eventPointer));
            switch (eventType)
            {
                case ConnectionEventConnected:
                    connection._log.Info("MsQuic QUIC handshake completed");
                    connection.RefreshRemoteEndPoint();
                    connection._connected.TrySetResult();
                    break;
                case ConnectionEventShutdownByTransport:
                case ConnectionEventShutdownByPeer:
                    connection._connected.TrySetException(new IOException("QUIC connection closed during handshake."));
                    break;
                case ConnectionEventShutdownComplete:
                    connection.OnShutdownComplete(handle);
                    break;
                case ConnectionEventPeerAddressChanged:
                {
                    nint address = Marshal.ReadIntPtr(eventPointer, IntPtr.Size);
                    IPEndPoint? endpoint = ReadEndpoint(address);
                    if (endpoint is not null)
                    {
                        connection.RemoteEndPoint = endpoint;
                        connection._peerAddressChanged?.Invoke(endpoint);
                    }

                    break;
                }
                case ConnectionEventDatagramStateChanged:
                    connection.OnDatagramStateChanged(eventPointer);
                    break;
                case ConnectionEventDatagramReceived:
                    connection.OnDatagramReceived(eventPointer);
                    break;
                case ConnectionEventDatagramSendStateChanged:
                    connection.OnDatagramSendStateChanged(eventPointer);
                    break;
                case ConnectionEventPeerCertificateReceived:
                    return connection.ValidatePeerCertificate(eventPointer);
            }

            return 0;
        }
        catch (Exception ex)
        {
            try
            {
                if (context != 0 && GCHandle.FromIntPtr(context).Target is MsQuicConnection connection)
                {
                    connection._log.Error("MsQuic connection callback failed", ("err", ex.Message));
                }
            }
            catch
            {
                // Keep exceptions inside the native callback boundary.
            }

            return OperatingSystem.IsWindows() ? 0x80004005u : 1u;
        }
    }

    private void OnDatagramStateChanged(nint eventPointer)
    {
        byte sendEnabled = Marshal.ReadByte(eventPointer, IntPtr.Size);
        int maxSendLength = (ushort)Marshal.ReadInt16(eventPointer, IntPtr.Size + 1);
        Volatile.Write(ref _maxDatagramBytes, maxSendLength);
        Volatile.Write(ref _sendEnabled, sendEnabled == 0 ? 0 : 1);
    }

    private void OnDatagramReceived(nint eventPointer)
    {
        nint bufferPointer = Marshal.ReadIntPtr(eventPointer, IntPtr.Size);
        if (bufferPointer == 0)
        {
            return;
        }

        var buffer = Marshal.PtrToStructure<MsQuicApi.QuicBuffer>(bufferPointer);
        if (buffer.Buffer != 0 && buffer.Length > 0 && buffer.Length <= 65535)
        {
            unsafe
            {
            DatagramHandler?.Invoke(this, new ReadOnlySpan<byte>((void*)buffer.Buffer, checked((int)buffer.Length)));
            }
        }
    }

    private void OnDatagramSendStateChanged(nint eventPointer)
    {
        nint context = Marshal.ReadIntPtr(eventPointer, IntPtr.Size);
        int state = Marshal.ReadInt32(eventPointer, IntPtr.Size * 2);
        if (state >= DatagramStateFinal)
        {
            ReleaseSend(context);
        }
    }

    private uint ValidatePeerCertificate(nint eventPointer)
    {
        nint certificate = Marshal.ReadIntPtr(eventPointer, IntPtr.Size);
        if (certificate == 0 || _expectedPeerSpki is null)
        {
            return 0;
        }

        byte[] der;
        if (OperatingSystem.IsWindows())
        {
            nint encoded = Marshal.ReadIntPtr(certificate, IntPtr.Size);
            int length = Marshal.ReadInt32(certificate, IntPtr.Size * 2);
            if (encoded == 0 || length <= 0 || length > 64 * 1024)
            {
                return 0x80004005u;
            }

            der = new byte[length];
            Marshal.Copy(encoded, der, 0, length);
        }
        else
        {
            var nativeBuffer = Marshal.PtrToStructure<MsQuicApi.QuicBuffer>(certificate);
            if (nativeBuffer.Buffer == 0 || nativeBuffer.Length == 0 || nativeBuffer.Length > 64 * 1024)
            {
                return 1;
            }

            der = new byte[checked((int)nativeBuffer.Length)];
            Marshal.Copy(nativeBuffer.Buffer, der, 0, der.Length);
        }

        using X509Certificate2 peerCertificate = X509CertificateLoader.LoadCertificate(der);
        bool valid = ExportSubjectPublicKeyInfo(peerCertificate).AsSpan().SequenceEqual(_expectedPeerSpki);
        return valid ? 0 : OperatingSystem.IsWindows() ? 0x80004005u : 1u;
    }

    private static byte[] ExportSubjectPublicKeyInfo(X509Certificate2 certificate)
    {
        using ECDsa? ecdsa = certificate.GetECDsaPublicKey();
        if (ecdsa is not null)
        {
            return ecdsa.ExportSubjectPublicKeyInfo();
        }

        using RSA? rsa = certificate.GetRSAPublicKey();
        if (rsa is not null)
        {
            return rsa.ExportSubjectPublicKeyInfo();
        }

        using DSA? dsa = certificate.GetDSAPublicKey();
        if (dsa is not null)
        {
            return dsa.ExportSubjectPublicKeyInfo();
        }

        return [];
    }

    private void ReleaseSend(nint context)
    {
        SendState? state;
        lock (_sendGate)
        {
            if (!_sends.Remove(context, out state))
            {
                return;
            }
        }

        if (GCHandle.FromIntPtr(context).IsAllocated)
        {
            GCHandle.FromIntPtr(context).Free();
        }

        state.Release();
    }

    private void OnShutdownComplete(nint handle)
    {
        if (Interlocked.Exchange(ref _handle, 0) == 0)
        {
            return;
        }

        MsQuicApi.ConnectionClose(handle);
        lock (_sendGate)
        {
            foreach ((nint context, SendState state) in _sends)
            {
                GCHandle.FromIntPtr(context).Free();
                state.Release();
            }

            _sends.Clear();
        }

        if (_configuration != 0)
        {
            MsQuicApi.ConfigurationClose(_configuration);
        }

        if (_registration != 0)
        {
            MsQuicApi.RegistrationClose(_registration);
        }

        _self.Free();
        _shutdown.TrySetResult();
        _closed?.Invoke(this);
    }

    private static IPEndPoint? ReadEndpoint(nint address)
    {
        if (address == 0)
        {
            return null;
        }

        ushort family = unchecked((ushort)Marshal.ReadInt16(address));
        if (family == 2)
        {
            byte[] bytes = new byte[4];
            Marshal.Copy(address + 4, bytes, 0, bytes.Length);
            int port = (Marshal.ReadByte(address, 2) << 8) | Marshal.ReadByte(address, 3);
            return new IPEndPoint(new IPAddress(bytes), port);
        }

        ushort ipv6Family = OperatingSystem.IsWindows() ? (ushort)23 : OperatingSystem.IsMacOS() ? (ushort)30 : (ushort)10;
        if (family == ipv6Family)
        {
            byte[] bytes = new byte[16];
            Marshal.Copy(address + 8, bytes, 0, bytes.Length);
            int port = (Marshal.ReadByte(address, 2) << 8) | Marshal.ReadByte(address, 3);
            long scope = unchecked((uint)Marshal.ReadInt32(address, 24));
            return new IPEndPoint(new IPAddress(bytes, scope), port);
        }

        return null;
    }

    private static IPEndPoint? ReadEndpoint(byte[] address)
    {
        unsafe
        {
            fixed (byte* addressBytes = address)
            {
                return ReadEndpoint((nint)addressBytes);
            }
        }
    }

    private static ushort ReadAddressFamily(byte[] address) => (ushort)(address[0] | (address[1] << 8));

    private sealed class SendState(byte[] buffer, GCHandle pin, nint nativeBuffer)
    {
        public void Release()
        {
            if (pin.IsAllocated)
            {
                pin.Free();
            }

            if (nativeBuffer != 0)
            {
                Marshal.FreeHGlobal(nativeBuffer);
            }

            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

}