namespace trojan_csharp.Models;

using System.Buffers;
using System.Buffers.Binary;
using System.Diagnostics.CodeAnalysis;
using System.Net;
using System.Net.Sockets;
using System.Text;
using Microsoft.AspNetCore.Connections;

public static class Trojan
{
    public static bool TryParseHeader(ReadOnlySpan<byte> buffer, [NotNullWhen(true)] out string? token, [NotNullWhen(true)] out int? command, [NotNullWhen(true)] out EndPoint? target, out int offset)
    {
        offset = 0;
        token = null;
        target = null;
        command = null;
        // TOKEN
        if (buffer.Length < 60) return false;
        token = Encoding.ASCII.GetString(buffer.Slice(offset, 56));
        offset += 56;
        // CRLF 2bit
        if (BinaryPrimitives.ReadUInt16BigEndian(buffer.Slice(offset, 2)) != 0x0D0A) {
            return false;
        }
        offset += 2;
        // COMMAND 1bit
        command = buffer[offset++];
        if (command is not (0x01 or 0x03)) return false;
        // ATYP 1bit
        var atyp = buffer[offset++];
        // ADDRESS
        ReadOnlySpan<byte> host;
        switch (atyp) {
            case 0x01 or 0x04:
                var addrLen = atyp * 4;
                if (offset + addrLen > buffer.Length) goto default;
                host = buffer.Slice(offset, addrLen);
                offset += addrLen;
                break;
            case 0x03 when offset + 1 < buffer.Length:
                var domainLen = buffer[offset++];
                if (offset + domainLen > buffer.Length) goto default;
                host = buffer.Slice(offset, domainLen);
                offset += domainLen;
                break;
            default: return false;
        }
        // PORT
        if (offset + 2 > buffer.Length) return false;
        var port = BinaryPrimitives.ReadUInt16BigEndian(buffer.Slice(offset, 2));
        offset += 2;
        // CRLF
        if (offset + 2 > buffer.Length) return false;
        if (BinaryPrimitives.ReadUInt16BigEndian(buffer.Slice(offset, 2)) != 0x0D0A) {
            return false;
        }
        offset += 2;
        target = atyp is 0x01 or 0x04 ? new IPEndPoint(new IPAddress(host), port) : new DnsEndPoint(Encoding.ASCII.GetString(host), port);
        return true;
    }

    public static class Udp
    {
        private static (int HeaderLength, int PayloadLengthOffset) UdpWriteHeader(IPEndPoint remoteEndPoint, Span<byte> buffer)
        {
            var address = remoteEndPoint.Address;
            var addrLen = address.AddressFamily == AddressFamily.InterNetwork ? 4 : 16;
            var headerLength = 1 + addrLen + 2 + 2 + 2;
            var offset = 0;
            // 1. ATYP
            buffer[offset++] = (byte)(addrLen == 4 ? 0x01 : 0x04);
            // 2. IP 地址
            address.TryWriteBytes(buffer[offset..], out _);
            offset += addrLen;
            // 3. 端口
            BinaryPrimitives.WriteUInt16BigEndian(buffer[offset..], (ushort)remoteEndPoint.Port);
            offset += 2;
            // 4. 预留长度字段位置
            var payloadLengthOffset = offset;
            offset += 2;
            // 5. CRLF
            BinaryPrimitives.WriteUInt16BigEndian(buffer[offset..], 0x0D0A);
            return (headerLength, payloadLengthOffset);
        }

        private static bool TryParseUdpPacket(ReadOnlySequence<byte> sequence, out int offset, out ushort payloadLength)
        {
            offset = 0;
            payloadLength = 0;
            if (sequence.IsEmpty) {
                return false;
            }
            var buffer = sequence.FirstSpan;
            if (buffer.Length < 7) {
                return false;
            }
            var addressLength = buffer[offset++] switch {
                0x01 => 4,
                0x03 => buffer[offset++],
                0x04 => 16,
                _ => -1
            };
            if (addressLength < 0 || offset + addressLength + 6 > buffer.Length) {
                return false;
            }
            offset += addressLength + 2;
            payloadLength = BinaryPrimitives.ReadUInt16BigEndian(buffer.Slice(offset, 2));
            offset += 4;
            return sequence.Length >= offset + payloadLength;
        }

        public static async Task CopyToAsync(ConnectionContext ctx, Socket socket, CancellationToken ct)
        {
            while (!ct.IsCancellationRequested) {
                var result = await ctx.Transport.Input.ReadAsync(ct);
                var sequence = result.Buffer;
                var consumed = sequence.Start;
                try {
                    if (TryParseUdpPacket(sequence, out var offset, out var payloadLength)) {
                        foreach (var seq in sequence.Slice(offset, payloadLength)) {
                            await socket.SendAsync(seq, SocketFlags.None, ct);
                        }
                        consumed = sequence.GetPosition(offset + payloadLength);
                    }
                    if (result.IsCompleted || result.IsCanceled) {
                        break;
                    }
                } finally {
                    ctx.Transport.Input.AdvanceTo(consumed, sequence.End);
                }
            }
        }

        public static async Task CopyToAsync(Socket socket, ConnectionContext ctx, CancellationToken ct)
        {
            using var memoryOwner = MemoryPool<byte>.Shared.Rent((64 * 1024) + 64);
            var buffer = memoryOwner.Memory;
            var span = buffer.Span;
            if (socket.RemoteEndPoint is IPEndPoint remoteEndPoint) {
                var (headerLength, payloadLengthOffset) = UdpWriteHeader(remoteEndPoint, span);
                while (!ct.IsCancellationRequested) {
                    var length = await socket.ReceiveAsync(buffer[headerLength..], SocketFlags.None, ct);
                    BinaryPrimitives.WriteUInt16BigEndian(buffer[payloadLengthOffset..].Span, (ushort)length);
                    var result = await ctx.Transport.Output.WriteAsync(buffer[..(headerLength + length)], ct);
                    if (result.IsCompleted || result.IsCanceled) {
                        break;
                    }
                }
            }
        }
    }
}