using System.Buffers.Binary;
using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;

namespace CampusNet.Net;

/// <summary>
/// Photon's ENet-style reliable UDP (wire details verified from a capture, docs/photon-wire-findings.md):
/// 12-byte packet header + N commands, per-channel reliable sequencing with acks and retransmission,
/// unreliable sequencing, fragmentation above the MTU, ping/keep-alive, server-time via ack timestamp.
/// One instance per listening port; the service on top receives whole messages per peer.
/// </summary>
public sealed class PhotonUdpServer
{
    public const int Mtu = 1200;
    private const int PacketHeader = 12, CommandHeader = 12;
    private const byte CmdAck = 1, CmdConnect = 2, CmdVerifyConnect = 3, CmdDisconnect = 4, CmdPing = 5,
                       CmdReliable = 6, CmdUnreliable = 7, CmdFragment = 8, CmdServerTime = 12;
    private const byte FlagReliable = 1, FlagUnsequenced = 2;
    private const byte ChannelSystem = 255;
    private const int PeerTimeoutMs = 12000, PingIntervalMs = 1000, MaxRetries = 8;

    public readonly int Port;
    public readonly string Name;
    private readonly Socket _sock;
    private readonly ConcurrentDictionary<EndPoint, Peer> _peers = new();
    private readonly System.Diagnostics.Stopwatch _clock = System.Diagnostics.Stopwatch.StartNew();
    private readonly Random _rng = new();
    private int ServerTime => (int)_clock.ElapsedMilliseconds;

    public event Action<Peer, byte[], bool>? OnMessage;   // a complete message (0xF3 ...), reliable?
    public event Action<Peer>? OnDisconnect;
    public event Action<Peer>? OnConnect;

    public PhotonUdpServer(string name, int port)
    {
        Name = name; Port = port;
        _sock = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
        try { _sock.IOControl(-1744830452, new byte[] { 0 }, null); } catch { }   // SIO_UDP_CONNRESET off (Windows)
        _sock.Bind(new IPEndPoint(IPAddress.Any, port));
    }

    public int PeerCount => _peers.Count;
    public IEnumerable<Peer> Peers => _peers.Values;

    // ------------------------------------------------------------ peer

    public sealed class Peer
    {
        public readonly EndPoint Remote;
        public readonly short PeerId;
        public int Challenge;
        public bool Connected;
        public long LastRecv, LastPing, LastSend;
        public int Rtt = 200;
        public object? Session;            // the service's per-connection state
        internal readonly Dictionary<byte, ChannelState> Channels = new();
        internal readonly List<PendingAck> AckQueue = new();
        internal readonly List<OutCommand> OutQueue = new();
        internal readonly Dictionary<long, OutCommand> UnackedBySeq = new();
        internal readonly Dictionary<long, Fragmenter> Fragments = new();
        internal readonly HashSet<byte> UnknownTypesLogged = new();
        internal int UnreliableOutSeq;
        internal bool Disconnecting;
        internal readonly TokenBucket InBytes = new(Guard.InBytesPerSec, Guard.InBytesBurst);
        public readonly string Ip;
        public Peer(EndPoint remote, short id) { Remote = remote; PeerId = id; Ip = Guard.IpOf(remote); }
        public override string ToString() => $"{Remote}#{PeerId}";
    }

    internal sealed class ChannelState
    {
        public int ReliableOut;       // last reliable seq sent
        public int ReliableInExpected = 1;
        public int UnreliableIn;
        public readonly SortedDictionary<int, byte[]> Pending = new();   // out-of-order reliable commands (raw command bytes)
    }

    internal struct PendingAck { public byte Channel; public int Seq; public int SentTime; }

    internal sealed class OutCommand
    {
        public byte Type, Channel, Flags; public int Seq; public byte[] Payload = Array.Empty<byte>();
        public long FirstSent, LastSent; public int Retries; public bool Reliable => (Flags & FlagReliable) != 0;
        public int Length => CommandHeader + Payload.Length;
    }

    internal sealed class Fragmenter { public int Count, Total; public readonly SortedDictionary<int, byte[]> Parts = new(); }

    // ------------------------------------------------------------ run loop

    public async Task RunAsync(CancellationToken ct)
    {
        Console.WriteLine($"[{Name}] listening on udp/{Port}");
        var buf = new byte[65536];
        var tick = Task.Run(() => TickLoop(ct), ct);
        while (!ct.IsCancellationRequested)
        {
            SocketReceiveFromResult r;
            try { r = await _sock.ReceiveFromAsync(buf, SocketFlags.None, new IPEndPoint(IPAddress.Any, 0), ct); }
            catch (OperationCanceledException) { break; }
            catch (SocketException) { continue; }
            catch (Exception ex) { Console.WriteLine($"[{Name}] recv error: {ex.Message}"); continue; }
            try { HandlePacket(r.RemoteEndPoint, buf.AsSpan(0, r.ReceivedBytes)); }
            catch (Exception ex) { Console.WriteLine($"[{Name}] packet from {r.RemoteEndPoint} failed: {ex}"); }
        }
        await tick;
    }

    private async Task TickLoop(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            try { await Task.Delay(20, ct); } catch { break; }
            long now = _clock.ElapsedMilliseconds;
            foreach (var p in _peers.Values)
            {
                try
                {
                    lock (p)
                    {
                        if (now - p.LastRecv > PeerTimeoutMs) { DropPeer(p, "timeout"); continue; }
                        if (p.Connected && now - p.LastPing > PingIntervalMs)
                        {
                            p.LastPing = now;
                            Enqueue(p, new OutCommand { Type = CmdPing, Channel = ChannelSystem, Flags = FlagReliable });
                        }
                        // Retransmit unacked reliable commands.
                        foreach (var c in p.UnackedBySeq.Values.ToList())
                        {
                            int rto = Math.Clamp(p.Rtt * 2 + 50, 200, 2000) << Math.Min(c.Retries, 3);
                            if (now - c.LastSent < rto) continue;
                            if (c.Retries++ >= MaxRetries) { DropPeer(p, "unacked command"); break; }
                            p.OutQueue.Add(c);
                        }
                        Flush(p);
                    }
                }
                catch (Exception ex) { Console.WriteLine($"[{Name}] tick {p}: {ex.Message}"); }
            }
        }
    }

    // ------------------------------------------------------------ receive

    private void HandlePacket(EndPoint from, ReadOnlySpan<byte> pkt)
    {
        if (pkt.Length < PacketHeader) return;
        short peerId = BinaryPrimitives.ReadInt16BigEndian(pkt);
        int count = pkt[3];
        int sentTime = BinaryPrimitives.ReadInt32BigEndian(pkt.Slice(4));
        int challenge = BinaryPrimitives.ReadInt32BigEndian(pkt.Slice(8));

        if (!_peers.TryGetValue(from, out var peer))
        {
            // Only a CONNECT may create a peer.
            if (pkt.Length < PacketHeader + CommandHeader || pkt[PacketHeader] != CmdConnect) return;
            string ip = Guard.IpOf(from);
            if (Guard.IsBanned(ip)) return;
            if (!Guard.AllowConnect(ip)) { Guard.BanIp(ip, $"more than {Guard.ConnectsPerIpPer10s} connects in 10s on {Name}"); return; }
            if (_peers.Values.Count(p => p.Ip == ip) >= Guard.MaxPeersPerIp)
            {
                Console.WriteLine($"[{Name}] connect from {from} refused: {Guard.MaxPeersPerIp} peers already from {ip}");
                return;
            }
            short id;
            do { id = (short)_rng.Next(1, short.MaxValue); } while (_peers.Values.Any(p => p.PeerId == id));
            peer = new Peer(from, id) { Challenge = challenge };
            _peers[from] = peer;
            Console.WriteLine($"[{Name}] connect from {from} -> peer {id}");
        }

        lock (peer)
        {
            long nowMs = _clock.ElapsedMilliseconds;
            peer.LastRecv = nowMs;
            // Inbound flood: over the byte budget the packet is dropped before any decode work;
            // a peer that stays over it is closed and its address cooled.
            if (!peer.InBytes.Take(pkt.Length, nowMs))
            {
                if (peer.InBytes.DroppedSince == 1) Console.WriteLine($"[{Name}] {peer} over {Guard.InBytesPerSec} B/s inbound - dropping packets");
                if (nowMs - peer.InBytes.ViolatingSinceMs > Guard.FloodKickAfterSec * 1000L)
                {
                    Guard.BanIp(peer.Ip, $"inbound flood on {Name} ({peer.InBytes.Dropped} packets dropped)");
                    DropPeer(peer, "inbound flood");
                }
                return;
            }
            int pos = PacketHeader;
            for (int i = 0; i < count && pos + CommandHeader <= pkt.Length; i++)
            {
                byte type = pkt[pos], channel = pkt[pos + 1], flags = pkt[pos + 2];
                int len = BinaryPrimitives.ReadInt32BigEndian(pkt.Slice(pos + 4));
                int seq = BinaryPrimitives.ReadInt32BigEndian(pkt.Slice(pos + 8));
                if (len < CommandHeader || pos + len > pkt.Length) break;
                var body = pkt.Slice(pos + CommandHeader, len - CommandHeader);
                HandleCommand(peer, type, channel, flags, seq, sentTime, body);
                pos += len;
            }
            Flush(peer);
        }
    }

    private void HandleCommand(Peer peer, byte type, byte channel, byte flags, int seq, int sentTime, ReadOnlySpan<byte> body)
    {
        if ((flags & FlagReliable) != 0 && type != CmdAck)
            peer.AckQueue.Add(new PendingAck { Channel = channel, Seq = seq, SentTime = sentTime });

        switch (type)
        {
            case CmdAck:
            {
                if (body.Length < 8) return;
                int ackedSeq = BinaryPrimitives.ReadInt32BigEndian(body);
                int ackedSent = BinaryPrimitives.ReadInt32BigEndian(body.Slice(4));
                if (peer.UnackedBySeq.Remove(Key(channel, ackedSeq), out var cmd))
                {
                    int rtt = ServerTime - ackedSent;
                    if (rtt >= 0 && rtt < 5000) peer.Rtt = (peer.Rtt * 7 + rtt) / 8;
                }
                return;
            }
            case CmdConnect:
            {
                // Answer VerifyConnect with the peer id we assigned; the rest echoes the client's connect block.
                var payload = new byte[32];
                if (body.Length >= 32) body.Slice(0, 32).CopyTo(payload);
                BinaryPrimitives.WriteInt16BigEndian(payload, peer.PeerId);
                for (int i = 12; i < 32; i++) payload[i] = 0;
                if (!peer.Connected)
                {
                    peer.Connected = true;
                    peer.LastPing = _clock.ElapsedMilliseconds;
                    Enqueue(peer, new OutCommand { Type = CmdVerifyConnect, Channel = ChannelSystem, Flags = FlagReliable, Payload = payload });
                    OnConnect?.Invoke(peer);
                }
                return;
            }
            case CmdDisconnect:
                DropPeer(peer, "client disconnect");
                return;
            case CmdPing:
            case CmdServerTime:
                return;   // the ack (with our header sentTime) is the whole answer
            case CmdReliable:
                DeliverReliable(peer, channel, seq, body.ToArray());
                return;
            case CmdUnreliable:
            {
                if (body.Length < 4) return;
                int useq = BinaryPrimitives.ReadInt32BigEndian(body);
                var ch = Channel(peer, channel);
                if (useq <= ch.UnreliableIn) return;   // late/duplicate
                ch.UnreliableIn = useq;
                Dispatch(peer, body.Slice(4).ToArray(), false);
                return;
            }
            case CmdFragment:
            {
                if (body.Length < 20) return;
                int start = BinaryPrimitives.ReadInt32BigEndian(body);
                int fcount = BinaryPrimitives.ReadInt32BigEndian(body.Slice(4));
                int total = BinaryPrimitives.ReadInt32BigEndian(body.Slice(12));
                int offset = BinaryPrimitives.ReadInt32BigEndian(body.Slice(16));
                long key = Key(channel, start);
                // The fragment header is client-supplied: bound the reassembly buffer and the number
                // of half-built messages before either becomes an allocation the client controls.
                if (total <= 0 || total > Guard.MaxFragmentTotal || fcount <= 0 || fcount > Guard.MaxFragmentTotal / 64 || offset < 0 || offset >= total)
                {
                    Guard.BanIp(peer.Ip, $"bad fragment header on {Name} (total {total}, count {fcount}, offset {offset})");
                    DropPeer(peer, "bad fragment");
                    return;
                }
                if (!peer.Fragments.TryGetValue(key, out var f))
                {
                    if (peer.Fragments.Count >= Guard.MaxFragmentSets) { DropPeer(peer, "too many partial fragments"); return; }
                    peer.Fragments[key] = f = new Fragmenter { Count = fcount, Total = total };
                }
                f.Parts[offset] = body.Slice(20).ToArray();
                if (f.Parts.Count < f.Count) return;
                peer.Fragments.Remove(key);
                var whole = new byte[f.Total]; int p = 0;
                foreach (var part in f.Parts.Values) { Buffer.BlockCopy(part, 0, whole, p, Math.Min(part.Length, whole.Length - p)); p += part.Length; }
                // Fragments occupy reliable seqs start..start+count-1; deliver in order at `start`.
                DeliverReliable(peer, channel, start, whole, fcount);
                return;
            }
            default:
                // Unknown command type: say so once per (peer, type) - a missing delivery mode
                // (unsequenced variants are 11/14/15/16) shows up here instead of vanishing.
                if (peer.UnknownTypesLogged.Add(type))
                    Console.WriteLine($"[{Name}] {peer} sent command type {type} (channel {channel}, flags {flags}, {body.Length} bytes) - not handled, dropped");
                return;
        }
    }

    private static long Key(byte channel, int seq) => ((long)channel << 32) | (uint)seq;
    private ChannelState Channel(Peer p, byte c) { if (!p.Channels.TryGetValue(c, out var s)) p.Channels[c] = s = new ChannelState(); return s; }

    private void DeliverReliable(Peer peer, byte channel, int seq, byte[] message, int span = 1)
    {
        var ch = Channel(peer, channel);
        if (seq < ch.ReliableInExpected) return;          // duplicate (already acked above)
        if (seq > ch.ReliableInExpected)
        {
            if (ch.Pending.Count >= Guard.MaxPendingReliable) { DropPeer(peer, "reliable window overflow"); return; }
            ch.Pending[seq] = message; return;
        }
        ch.ReliableInExpected = seq + span;
        Dispatch(peer, message, true);
        while (ch.Pending.Count > 0)
        {
            var first = ch.Pending.First();
            if (first.Key != ch.ReliableInExpected) break;
            ch.Pending.Remove(first.Key);
            ch.ReliableInExpected++;
            Dispatch(peer, first.Value, true);
        }
    }

    private void Dispatch(Peer peer, byte[] message, bool reliable)
    {
        if (message.Length == 0) return;
        try { OnMessage?.Invoke(peer, message, reliable); }
        catch (Exception ex) { Console.WriteLine($"[{Name}] handler error for {peer}: {ex}"); }
    }

    // ------------------------------------------------------------ send

    public void SendReliable(Peer peer, byte[] message, byte channel = 0)
    {
        lock (peer)
        {
            if (peer.Disconnecting) return;
            int maxBody = Mtu - PacketHeader - CommandHeader;
            if (message.Length <= maxBody)
            {
                Enqueue(peer, new OutCommand { Type = CmdReliable, Channel = channel, Flags = FlagReliable, Payload = message });
                return;
            }
            int chunk = maxBody - 20;
            int count = (message.Length + chunk - 1) / chunk;
            var ch = Channel(peer, channel);
            int start = ch.ReliableOut + 1;
            for (int i = 0; i < count; i++)
            {
                int off = i * chunk, len = Math.Min(chunk, message.Length - off);
                var payload = new byte[20 + len];
                BinaryPrimitives.WriteInt32BigEndian(payload, start);
                BinaryPrimitives.WriteInt32BigEndian(payload.AsSpan(4), count);
                BinaryPrimitives.WriteInt32BigEndian(payload.AsSpan(8), i);
                BinaryPrimitives.WriteInt32BigEndian(payload.AsSpan(12), message.Length);
                BinaryPrimitives.WriteInt32BigEndian(payload.AsSpan(16), off);
                Buffer.BlockCopy(message, off, payload, 20, len);
                Enqueue(peer, new OutCommand { Type = CmdFragment, Channel = channel, Flags = FlagReliable, Payload = payload });
            }
        }
    }

    public void SendUnreliable(Peer peer, byte[] message, byte channel = 0)
    {
        lock (peer)
        {
            if (peer.Disconnecting) return;
            if (message.Length > Mtu - PacketHeader - CommandHeader - 4) { SendReliable(peer, message, channel); return; }
            var payload = new byte[4 + message.Length];
            BinaryPrimitives.WriteInt32BigEndian(payload, ++peer.UnreliableOutSeq);
            Buffer.BlockCopy(message, 0, payload, 4, message.Length);
            Enqueue(peer, new OutCommand { Type = CmdUnreliable, Channel = channel, Flags = 0, Payload = payload });
        }
    }

    public void Disconnect(Peer peer, string why)
    {
        lock (peer)
        {
            if (peer.Disconnecting) return;
            Enqueue(peer, new OutCommand { Type = CmdDisconnect, Channel = ChannelSystem, Flags = FlagUnsequenced });
            Flush(peer);
            DropPeer(peer, why);
        }
    }

    private void Enqueue(Peer peer, OutCommand c)
    {
        var ch = Channel(peer, c.Channel);
        if (c.Reliable)
        {
            c.Seq = ++ch.ReliableOut;
            peer.UnackedBySeq[Key(c.Channel, c.Seq)] = c;
        }
        else
        {
            // ENet sequences unreliable commands AGAINST the channel's reliable stream: the header's
            // sequence field must carry the last reliable seq sent on the channel. The receiver
            // holds an unreliable command until it has that reliable command, and DISCARDS one whose
            // reliable seq is older than what it already has. Stamping 0 here meant every unreliable
            // event the server relayed - all the avatar movement - was thrown away as stale the
            // moment the client had received its first reliable message (the join response).
            c.Seq = ch.ReliableOut;
        }
        peer.OutQueue.Add(c);
    }

    /// <summary>Pack queued acks + commands into MTU-sized packets and send them.</summary>
    private void Flush(Peer peer)
    {
        if (peer.AckQueue.Count == 0 && peer.OutQueue.Count == 0) return;
        long now = _clock.ElapsedMilliseconds;
        var pkt = new byte[Mtu];
        int pos = PacketHeader, n = 0;

        void Send()
        {
            if (n == 0) return;
            BinaryPrimitives.WriteInt16BigEndian(pkt, 0);
            pkt[2] = 0; pkt[3] = (byte)n;
            BinaryPrimitives.WriteInt32BigEndian(pkt.AsSpan(4), ServerTime);
            BinaryPrimitives.WriteInt32BigEndian(pkt.AsSpan(8), peer.Challenge);
            try { _sock.SendTo(pkt, 0, pos, SocketFlags.None, peer.Remote); } catch (Exception ex) { Console.WriteLine($"[{Name}] send to {peer} failed: {ex.Message}"); }
            pos = PacketHeader; n = 0;
        }
        void Put(byte type, byte channel, byte flags, int seq, ReadOnlySpan<byte> payload)
        {
            int len = CommandHeader + payload.Length;
            if (pos + len > Mtu) Send();
            pkt[pos] = type; pkt[pos + 1] = channel; pkt[pos + 2] = flags; pkt[pos + 3] = 4;
            BinaryPrimitives.WriteInt32BigEndian(pkt.AsSpan(pos + 4), len);
            BinaryPrimitives.WriteInt32BigEndian(pkt.AsSpan(pos + 8), seq);
            payload.CopyTo(pkt.AsSpan(pos + CommandHeader));
            pos += len; n++;
            if (n == 255) Send();
        }

        foreach (var a in peer.AckQueue)
        {
            Span<byte> ap = stackalloc byte[8];
            BinaryPrimitives.WriteInt32BigEndian(ap, a.Seq);
            BinaryPrimitives.WriteInt32BigEndian(ap.Slice(4), a.SentTime);
            Put(CmdAck, a.Channel, 0, 0, ap);
        }
        peer.AckQueue.Clear();
        foreach (var c in peer.OutQueue)
        {
            if (c.Reliable) { if (c.FirstSent == 0) c.FirstSent = now; c.LastSent = now; }
            Put(c.Type, c.Channel, c.Flags, c.Seq, c.Payload);
        }
        peer.OutQueue.Clear();
        Send();
        peer.LastSend = now;
    }

    /// <summary>Close a peer for abuse and cool its address (Guard.IpBanSec).</summary>
    public void Kick(Peer peer, string why, int? banSeconds = null)
    {
        Guard.BanIp(peer.Ip, $"{why} on {Name}", banSeconds);
        lock (peer) Disconnect(peer, why);
    }

    private void DropPeer(Peer peer, string why)
    {
        if (peer.Disconnecting) return;
        peer.Disconnecting = true;
        _peers.TryRemove(peer.Remote, out _);
        Console.WriteLine($"[{Name}] {peer} gone ({why})");
        try { OnDisconnect?.Invoke(peer); } catch (Exception ex) { Console.WriteLine($"[{Name}] disconnect handler: {ex}"); }
    }
}