using System.Security.Cryptography;
using Prospect.Unreal.Serialization;
using Serilog;

namespace Prospect.Unreal.Net;

/// <summary>
/// Chiffrement AES-256-GCM keyé directement par la PSK (pendant de l'AESGCMHandlerComponent
/// d'UE / flux SetEncryptionData+EnableEncryption). Hypothèse retenue après test live : le
/// client n'entame pas de handshake DTLS ; après NMT_EncryptionAck il attend une voie serveur
/// chiffrée. Pas de handshake → la PSK 32 o EST la clé AES.
///
/// Format paquet UE : [IV 12o][ciphertext][tag 16o].
///
/// ⚠️ Première implémentation — l'alignement bit exact du framing UE est à valider EN LIVE.
/// </summary>
public sealed class AesGcmHandlerComponent : HandlerComponent
{
    private static readonly ILogger Logger = Log.ForContext<AesGcmHandlerComponent>();

    private const int IvSize = 12;
    private const int TagSize = 16;

    private AesGcm? _aes;

    public AesGcmHandlerComponent(PacketHandler handler) : base(handler, nameof(AesGcmHandlerComponent))
    {
    }

    public override void Initialize()
    {
        SetActive(false);
        Initialized();
    }

    public override bool IsValid() => true;

    /// <summary>Active le chiffrement avec la clé (PSK 32 o). À appeler APRÈS avoir envoyé
    /// l'EncryptionAck en clair.</summary>
    public void Activate(byte[] key)
    {
        _aes = new AesGcm(key, TagSize);
        SetActive(true);
        Logger.Information("[AES] Chiffrement AES-256-GCM activé (clé {N} o)", key.Length);
    }

    public override void Incoming(FBitReader packet)
    {
        if (!IsActive() || _aes == null)
        {
            return;
        }

        var buf = ReadAlignedBytes(packet);
        if (buf.Length < IvSize + TagSize)
        {
            Logger.Warning("[AES] Incoming trop court ({N} o) — ignoré", buf.Length);
            packet.SetData(Array.Empty<byte>(), 0);
            return;
        }

        var iv = buf.AsSpan(0, IvSize);
        var tag = buf.AsSpan(buf.Length - TagSize, TagSize);
        var ct = buf.AsSpan(IvSize, buf.Length - IvSize - TagSize);
        var plain = new byte[ct.Length];
        try
        {
            _aes.Decrypt(iv, ct, tag, plain);
            Logger.Information("[AES] Incoming déchiffré {N}o -> {P}o head={H}", buf.Length, plain.Length,
                Convert.ToHexString(plain.AsSpan(0, Math.Min(plain.Length, 12))));
            packet.SetData(plain, (long)plain.Length * 8);
        }
        catch (CryptographicException ex)
        {
            Logger.Warning("[AES] Déchiffrement échoué ({M}) — paquet {N}o head={H}", ex.Message, buf.Length,
                Convert.ToHexString(buf.AsSpan(0, Math.Min(buf.Length, 16))));
            packet.SetError();
        }
    }

    public override void Outgoing(ref FBitWriter packet, FOutPacketTraits traits)
    {
        if (!IsActive() || _aes == null)
        {
            return;
        }

        var plain = packet.GetData();
        var nbytes = (int)packet.GetNumBytes();

        var iv = new byte[IvSize];
        RandomNumberGenerator.Fill(iv);
        var ct = new byte[nbytes];
        var tag = new byte[TagSize];
        _aes.Encrypt(iv, plain.AsSpan(0, nbytes), ct, tag);

        var outBuf = new byte[IvSize + nbytes + TagSize];
        iv.CopyTo(outBuf, 0);
        ct.CopyTo(outBuf, IvSize);
        tag.CopyTo(outBuf, IvSize + nbytes);

        var newPacket = new FBitWriter((long)outBuf.Length * 8 + 1, true, false);
        newPacket.SerializeBits(outBuf, outBuf.Length * 8);
        packet = newPacket;
        Logger.Information("[AES] Outgoing chiffré {P}o -> {N}o", nbytes, outBuf.Length);
    }

    public override int GetReservedPacketBits() => (IvSize + TagSize) * 8;

    public bool IsEnabled => IsActive();

    private static byte[] ReadAlignedBytes(FBitReader packet)
    {
        var bytes = packet.GetBytesLeft();
        if (bytes <= 0)
        {
            return Array.Empty<byte>();
        }

        var buf = new byte[bytes];
        unsafe
        {
            fixed (byte* p = buf)
            {
                packet.Serialize(p, bytes);
            }
        }
        return buf;
    }
}
