using System.IO.Compression;
using System.Numerics;
using System.Security.Cryptography;
using System.Text;

namespace MxAsbClient;

internal sealed class AsbSystemAuthenticator
{
    private static readonly byte[] PasswordSalt = Encoding.ASCII.GetBytes("ArchestrAService");
    private readonly BigInteger dhPrime;
    private readonly BigInteger dhGenerator;
    private readonly string hashAlgorithm;
    private readonly int keySize;
    private readonly byte[] solutionPassphrase;
    private readonly byte[] privateKey;
    private readonly byte[] localPublicKey;
    private byte[] remotePublicKey = [];
    private ulong nextMessageNumber = 1;
    /// Trace callback for the F28 canonical-XML reconciliation pass —
    /// when set, `Sign` dumps the request type, the UTF-8 bytes of
    /// `request.ToXml()`, the resulting HMAC, and the encrypted MAC +
    /// IV. Used by `MxAsbClient.Probe --dump-signed-xml` and ad-hoc
    /// live runs to capture the exact bytes the server's HMAC verifier
    /// recomputes against; the Rust port's `xml_canonical` emitter must
    /// produce byte-identical XML for the HMAC to round-trip.
    private readonly Action<string>? sharedTrace;

    public AsbSystemAuthenticator(string passphrase, AsbSolutionCryptoParameters cryptoParameters, Action<string>? trace = null)
    {
        sharedTrace = trace;
        dhPrime = cryptoParameters.Prime;
        dhGenerator = cryptoParameters.Generator;
        hashAlgorithm = cryptoParameters.HashAlgorithm;
        keySize = cryptoParameters.KeySize;
        trace?.Invoke("asb.stage=authenticator-passphrase-bytes");
        solutionPassphrase = Encoding.UTF8.GetBytes(passphrase);
        trace?.Invoke("asb.stage=authenticator-create-private");
        BigInteger privateKeyValue = CreatePrivateKey();
        trace?.Invoke("asb.stage=authenticator-private-ready");
        privateKey = privateKeyValue.ToByteArray();
        trace?.Invoke("asb.stage=authenticator-modpow");
        localPublicKey = BigInteger.ModPow(dhGenerator, privateKeyValue, dhPrime).ToByteArray();
        trace?.Invoke("asb.stage=authenticator-public-ready");
        ConnectionId = Guid.NewGuid();
    }

    public Guid ConnectionId { get; }

    public byte[] LocalPublicKey => localPublicKey;

    public bool UseApolloSigning { get; private set; }

    public void AcceptConnectResponse(ConnectResponse response)
    {
        remotePublicKey = response.ServicePublicKey?.Data ?? throw new InvalidOperationException("ASB connect response did not contain a service public key.");
        UseApolloSigning = response.ConnectionLifetime?.Contains(":V2", StringComparison.OrdinalIgnoreCase) == true;
    }

    public AuthenticationData CreateAuthenticationData()
    {
        byte[] clear = [.. localPublicKey, .. remotePublicKey];
        byte[] encrypted = Encrypt(clear, out byte[] iv);
        return new AuthenticationData
        {
            Data = encrypted,
            InitializationVector = iv,
        };
    }

    public void Sign(ConnectedRequest request, bool forceHmac = false)
    {
        ConnectionValidator validator = new()
        {
            ConnectionId = ConnectionId,
            MessageNumber = nextMessageNumber++,
            MessageAuthenticationCode = [],
            SignatureInitializationVector = [],
        };

        request.ConnectionValidator = validator;
        using HMAC? hmac = CreateHmac(forceHmac);
        if (hmac is null)
        {
            return;
        }

        string xmlText = request.ToXml();
        byte[] xmlBytes = Encoding.UTF8.GetBytes(xmlText);
        sharedTrace?.Invoke($"asb.sign.type={request.GetType().Name}");
        sharedTrace?.Invoke($"asb.sign.xml-utf8-len={xmlBytes.Length}");
        sharedTrace?.Invoke($"asb.sign.xml-b64={Convert.ToBase64String(xmlBytes)}");
        byte[] hash = hmac.ComputeHash(xmlBytes);
        sharedTrace?.Invoke($"asb.sign.hmac-b64={Convert.ToBase64String(hash)}");
        validator.MessageAuthenticationCode = Encrypt(hash, out byte[] iv);
        validator.SignatureInitializationVector = iv;
        sharedTrace?.Invoke($"asb.sign.encrypted-mac-b64={Convert.ToBase64String(validator.MessageAuthenticationCode)}");
        sharedTrace?.Invoke($"asb.sign.iv-b64={Convert.ToBase64String(iv)}");
    }

    private HMAC? CreateHmac(bool forceHmac)
    {
        return hashAlgorithm.ToLowerInvariant() switch
        {
            "md5" => new HMACMD5(CryptoKey),
            "sha1" => new HMACSHA1(CryptoKey),
            "sha512" => new HMACSHA512(CryptoKey),
            _ => forceHmac ? new HMACSHA1(CryptoKey) : null,
        };
    }

    private byte[] Encrypt(byte[] clear, out byte[] iv)
    {
        if (UseApolloSigning)
        {
            return EncryptApollo(clear, out iv);
        }

        return EncryptBaktun(clear, out iv);
    }

    private byte[] EncryptApollo(byte[] clear, out byte[] iv)
    {
        using Aes aes = Aes.Create();
        aes.Key = DeriveAesKey();
        iv = aes.IV;
        using MemoryStream output = new();
        using (CryptoStream crypto = new(output, aes.CreateEncryptor(), CryptoStreamMode.Write))
        {
            crypto.Write(clear, 0, clear.Length);
        }

        return output.ToArray();
    }

    private byte[] EncryptBaktun(byte[] clear, out byte[] iv)
    {
        using Aes aes = Aes.Create();
        aes.Key = DeriveAesKey();
        iv = aes.IV;
        using MemoryStream output = new();
        using (CryptoStream crypto = new(output, aes.CreateEncryptor(), CryptoStreamMode.Write))
        {
            using DeflateStream deflate = new(crypto, CompressionMode.Compress);
            deflate.Write(clear, 0, clear.Length);
        }

        return output.ToArray();
    }

    private byte[] DeriveAesKey()
    {
        byte[] cryptoKey = CryptoKey;
        byte[] aesKey = Rfc2898DeriveBytes.Pbkdf2(
            Convert.ToBase64String(cryptoKey),
            PasswordSalt,
            iterations: 1000,
            HashAlgorithmName.SHA1,
            outputLength: 16);
        sharedTrace?.Invoke($"asb.derive.crypto_key.len={cryptoKey.Length}");
        sharedTrace?.Invoke($"asb.derive.crypto_key.hex={Convert.ToHexString(cryptoKey)}");
        sharedTrace?.Invoke($"asb.derive.crypto_key.b64={Convert.ToBase64String(cryptoKey)}");
        sharedTrace?.Invoke($"asb.derive.aes_key.hex={Convert.ToHexString(aesKey)}");
        return aesKey;
    }

    private byte[] CryptoKey
    {
        get
        {
            byte[] shared = BigInteger.ModPow(new BigInteger(remotePublicKey), new BigInteger(privateKey), dhPrime).ToByteArray();
            return [.. shared, .. solutionPassphrase];
        }
    }

    private BigInteger CreatePrivateKey()
    {
        byte[] bytes = new byte[(keySize / 8) + 1];
        BigInteger value;
        do
        {
            RandomNumberGenerator.Fill(bytes);
            bytes[^1] = 0;
            value = new BigInteger(bytes);
        }
        while (value <= BigInteger.Zero || value >= dhPrime - BigInteger.One);

        return value;
    }
}