using System.Buffers.Binary;
using System.Net;
using System.Text;
using LumaTunnel.Shared.Models;

namespace LumaTunnel.Client.Core.Proxy;

internal static class Socks5Protocol
{
    public static async Task<ProxyTarget?> NegotiateAsync(Stream stream, CancellationToken cancellationToken)
    {
        var greeting = new byte[2];
        await ReadExactlyAsync(stream, greeting, cancellationToken).ConfigureAwait(false);
        if (greeting[0] != 5 || greeting[1] == 0)
            return null;
        var methods = new byte[greeting[1]];
        await ReadExactlyAsync(stream, methods, cancellationToken).ConfigureAwait(false);
        if (!methods.Contains((byte)0))
        {
            await stream.WriteAsync(new byte[] { 5, 0xFF }, cancellationToken).ConfigureAwait(false);
            return null;
        }
        await stream.WriteAsync(new byte[] { 5, 0 }, cancellationToken).ConfigureAwait(false);

        var header = new byte[4];
        await ReadExactlyAsync(stream, header, cancellationToken).ConfigureAwait(false);
        if (header[0] != 5 || header[1] != 1 || header[2] != 0)
        {
            await WriteReplyAsync(stream, header[1] == 1 ? (byte)0x01 : (byte)0x07, null, cancellationToken).ConfigureAwait(false);
            return null;
        }

        string host;
        IPAddress? address = null;
        switch (header[3])
        {
            case 1:
                var ipv4 = new byte[4];
                await ReadExactlyAsync(stream, ipv4, cancellationToken).ConfigureAwait(false);
                address = new IPAddress(ipv4);
                host = address.ToString();
                break;
            case 3:
                var length = stream.ReadByte();
                if (length <= 0)
                    return null;
                var domain = new byte[length];
                await ReadExactlyAsync(stream, domain, cancellationToken).ConfigureAwait(false);
                host = Encoding.ASCII.GetString(domain);
                break;
            case 4:
                var ipv6 = new byte[16];
                await ReadExactlyAsync(stream, ipv6, cancellationToken).ConfigureAwait(false);
                address = new IPAddress(ipv6);
                host = address.ToString();
                break;
            default:
                await WriteReplyAsync(stream, 0x08, null, cancellationToken).ConfigureAwait(false);
                return null;
        }

        var portBytes = new byte[2];
        await ReadExactlyAsync(stream, portBytes, cancellationToken).ConfigureAwait(false);
        var port = BinaryPrimitives.ReadUInt16BigEndian(portBytes);
        return port == 0 ? null : new ProxyTarget(host, port, address);
    }

    public static async Task WriteReplyAsync(Stream stream, byte reply, IPEndPoint? endpoint, CancellationToken cancellationToken)
    {
        endpoint ??= new IPEndPoint(IPAddress.Any, 0);
        var address = endpoint.Address;
        var addressBytes = address.GetAddressBytes();
        var response = new byte[4 + addressBytes.Length + 2];
        response[0] = 5;
        response[1] = reply;
        response[2] = 0;
        response[3] = address.AddressFamily == System.Net.Sockets.AddressFamily.InterNetwork ? (byte)1 : (byte)4;
        addressBytes.CopyTo(response, 4);
        BinaryPrimitives.WriteUInt16BigEndian(response.AsSpan(4 + addressBytes.Length), checked((ushort)endpoint.Port));
        await stream.WriteAsync(response, cancellationToken).ConfigureAwait(false);
    }

    private static async Task ReadExactlyAsync(Stream stream, Memory<byte> buffer, CancellationToken cancellationToken)
    {
        var offset = 0;
        while (offset < buffer.Length)
        {
            var count = await stream.ReadAsync(buffer[offset..], cancellationToken).ConfigureAwait(false);
            if (count == 0)
                throw new EndOfStreamException();
            offset += count;
        }
    }
}