namespace gatuna;

using System.Text;

static class Proto
{
    public const ushort EtherType = 0x6969;
    public const int EthHeaderLen = 14;
    public const byte Version = 2;
    public const int HeaderLen = 8;
    public const int DataHeaderLen = 16; // 8 common + 4 seq + 4 ack_seq
    public const int MaxPayload = 1480;

    public const byte TypeDiscover = 0x01;
    public const byte TypeManifest = 0x02;
    public const byte TypeOpen = 0x03;
    public const byte TypeOpenAck = 0x04;
    public const byte TypeOpenNak = 0x05;
    public const byte TypeData = 0x06;
    public const byte TypeClose = 0x07;
    public const byte TypePing = 0x0B;
    public const byte TypePong = 0x0C;

    public const byte ProtoTcp = 1;
    public const byte ProtoUdp = 2;

    public const byte ReasonUnspecified = 0;
    public const byte ReasonUnknownUpstream = 1;
    public const byte ReasonConnectFailed = 2;
    public const byte ReasonOversize = 3;
    public const byte ReasonUnknownSession = 4;
    public const byte ReasonMaxRetries = 5;
}

readonly record struct UpstreamEntry(
    byte Id,
    byte Protocol,
    ushort Port,
    string? Label)
{
    public string ProtoName => Protocol == Proto.ProtoTcp ? "tcp" : "udp";
}

abstract record Frame
{
    internal record Discover : Frame;
    internal record Manifest(string Hostname, UpstreamEntry[] Entries) : Frame;
    internal record Open(byte UpstreamId, byte Proto) : Frame;
    internal record OpenAck(uint SessionId, byte UpstreamId, byte Proto) : Frame;
    internal record OpenNak(byte UpstreamId, byte Reason) : Frame;
    internal record Data(uint SessionId, uint Seq, uint AckSeq, byte[] Payload) : Frame;
    internal record Close(uint SessionId, byte? Reason) : Frame;
    internal record Ping(ulong Nonce) : Frame;
    internal record Pong(ulong Nonce) : Frame;
}

static class FrameCodec
{
    /// Build a non-DATA frame: 8-byte common header + payload.
    static byte[] Build(byte type, uint sessionId, byte[] payload)
    {
        var buf = new byte[Proto.HeaderLen + payload.Length];
        buf[0] = Proto.Version;
        buf[1] = type;
        buf[2] = (byte)(sessionId >> 24);
        buf[3] = (byte)(sessionId >> 16);
        buf[4] = (byte)(sessionId >> 8);
        buf[5] = (byte)(sessionId & 0xFF);
        buf[6] = (byte)(payload.Length >> 8);
        buf[7] = (byte)(payload.Length & 0xFF);
        Buffer.BlockCopy(payload, 0, buf, Proto.HeaderLen, payload.Length);
        return buf;
    }

    /// Build a DATA frame: 8-byte common header + seq + ack_seq + payload.
    /// payload_len counts only the raw bytes, not seq/ack_seq.
    static byte[] BuildData(uint sessionId, uint seq, uint ackSeq, byte[] payload)
    {
        var buf = new byte[Proto.DataHeaderLen + payload.Length];
        buf[0] = Proto.Version;
        buf[1] = Proto.TypeData;
        buf[2] = (byte)(sessionId >> 24);
        buf[3] = (byte)(sessionId >> 16);
        buf[4] = (byte)(sessionId >> 8);
        buf[5] = (byte)(sessionId & 0xFF);
        buf[6] = (byte)(payload.Length >> 8);
        buf[7] = (byte)(payload.Length & 0xFF);
        buf[8] = (byte)(seq >> 24);
        buf[9] = (byte)(seq >> 16);
        buf[10] = (byte)(seq >> 8);
        buf[11] = (byte)(seq & 0xFF);
        buf[12] = (byte)(ackSeq >> 24);
        buf[13] = (byte)(ackSeq >> 16);
        buf[14] = (byte)(ackSeq >> 8);
        buf[15] = (byte)(ackSeq & 0xFF);
        Buffer.BlockCopy(payload, 0, buf, Proto.DataHeaderLen, payload.Length);
        return buf;
    }

    public static byte[] Encode(Frame frame)
    {
        return frame switch
        {
            Frame.Discover =>
                Build(Proto.TypeDiscover, 0, []),
            Frame.Manifest manifest =>
                Build(Proto.TypeManifest, 0, BuildManifestPayload(manifest.Hostname, manifest.Entries)),
            Frame.Open open =>
                Build(Proto.TypeOpen, 0, [open.UpstreamId, open.Proto]),
            Frame.OpenAck ack =>
                Build(Proto.TypeOpenAck, ack.SessionId, [ack.UpstreamId, ack.Proto]),
            Frame.OpenNak nak =>
                Build(Proto.TypeOpenNak, 0, [nak.UpstreamId, nak.Reason]),
            Frame.Data data =>
                BuildData(data.SessionId, data.Seq, data.AckSeq, data.Payload),
            Frame.Close close =>
                Build(Proto.TypeClose, close.SessionId,
                    close.Reason.HasValue ? [close.Reason.Value] : []),
            Frame.Ping ping =>
                Build(Proto.TypePing, 0, EncodeNonce(ping.Nonce)),
            Frame.Pong pong =>
                Build(Proto.TypePong, 0, EncodeNonce(pong.Nonce)),
            _ => throw new InvalidOperationException($"unknown frame type: {frame.GetType()}"),
        };
    }

    static byte[] BuildManifestPayload(string hostname, UpstreamEntry[] entries)
    {
        using var ms = new MemoryStream();
        var hnBytes = Encoding.UTF8.GetBytes(hostname);
        var hnLen = (byte)Math.Min(hnBytes.Length, 255);
        ms.WriteByte(hnLen);
        if (hnLen > 0)
            ms.Write(hnBytes, 0, hnLen);
        foreach (var e in entries)
        {
            var labelBytes = Encoding.UTF8.GetBytes(e.Label ?? "");
            var labelLen = (byte)Math.Min(labelBytes.Length, 255);
            ms.WriteByte(e.Id);
            ms.WriteByte(e.Protocol);
            ms.WriteByte((byte)(e.Port >> 8));
            ms.WriteByte((byte)(e.Port & 0xFF));
            ms.WriteByte(labelLen);
            if (labelLen > 0)
                ms.Write(labelBytes, 0, labelLen);
        }
        return ms.ToArray();
    }

    public static Frame? Parse(ReadOnlySpan<byte> buf)
    {
        if (buf.Length < Proto.HeaderLen)
            return null;
        if (buf[0] != Proto.Version)
            return null;
        var type = buf[1];
        var sessionId = (uint)(buf[2] << 24 | buf[3] << 16 | buf[4] << 8 | buf[5]);
        var payloadLen = (ushort)(buf[6] << 8 | buf[7]);

        // DATA frames have seq + ack_seq after the common header.
        if (type == Proto.TypeData)
        {
            if (buf.Length < Proto.DataHeaderLen + payloadLen)
                return null;
            var seq = (uint)(buf[8] << 24 | buf[9] << 16 | buf[10] << 8 | buf[11]);
            var ackSeq = (uint)(buf[12] << 24 | buf[13] << 16 | buf[14] << 8 | buf[15]);
            var payload = buf.Slice(Proto.DataHeaderLen, payloadLen);
            if (payload.Length > Proto.MaxPayload)
                return null;
            return new Frame.Data(sessionId, seq, ackSeq, payload.ToArray());
        }

        if (buf.Length < Proto.HeaderLen + payloadLen)
            return null;
        var payload2 = buf.Slice(Proto.HeaderLen, payloadLen);

        return type switch
        {
            Proto.TypeManifest => ParseManifest(payload2),
            Proto.TypeOpenAck when payload2.Length == 2 =>
                new Frame.OpenAck(sessionId, payload2[0], payload2[1]),
            Proto.TypeOpenNak when payload2.Length == 2 =>
                new Frame.OpenNak(payload2[0], payload2[1]),
            Proto.TypeClose when payload2.Length is 0 or 1 =>
                new Frame.Close(sessionId, payload2.Length == 1 ? payload2[0] : null),
            Proto.TypePong when payload2.Length == 8 =>
                new Frame.Pong(ParseNonce(payload2)),
            Proto.TypePing when payload2.Length == 8 =>
                new Frame.Ping(ParseNonce(payload2)),
            Proto.TypeDiscover when payload2.Length == 0 =>
                new Frame.Discover(),
            _ => null,
        };
    }

    static ulong ParseNonce(ReadOnlySpan<byte> payload)
    {
        ulong nonce = 0;
        for (int i = 0; i < 8; i++)
            nonce = (nonce << 8) | payload[i];
        return nonce;
    }

    static byte[] EncodeNonce(ulong nonce)
    {
        return [
            (byte)(nonce >> 56), (byte)(nonce >> 48),
            (byte)(nonce >> 40), (byte)(nonce >> 32),
            (byte)(nonce >> 24), (byte)(nonce >> 16),
            (byte)(nonce >> 8),  (byte)(nonce & 0xFF),
        ];
    }

    static Frame.Manifest? ParseManifest(ReadOnlySpan<byte> payload)
    {
        if (payload.Length < 1)
            return null;
        var hnLen = payload[0];
        if (1 + hnLen > payload.Length)
            return null;
        var hostname = hnLen == 0
            ? ""
            : Encoding.UTF8.GetString(payload[1..(1 + hnLen)]);
        var rest = payload[(1 + hnLen)..];

        var entries = new List<UpstreamEntry>();
        int i = 0;
        while (i < rest.Length)
        {
            if (i + 5 > rest.Length)
                return null;
            var id = rest[i];
            var proto = rest[i + 1];
            var port = (ushort)(rest[i + 2] << 8 | rest[i + 3]);
            var labelLen = rest[i + 4];
            i += 5;
            if (i + labelLen > rest.Length)
                return null;
            string? label = labelLen == 0
                ? null
                : Encoding.UTF8.GetString(rest[i..(i + labelLen)]);
            i += labelLen;
            entries.Add(new UpstreamEntry(id, proto, port, label));
        }
        return new Frame.Manifest(hostname, entries.ToArray());
    }
}