namespace trojan_csharp;

using System.Net;
using System.Net.Sockets;
using Microsoft.AspNetCore.Connections;
using Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets;
using Models;
using Models.Interfaces;

public class TrojanConnectionHandler(SocketConnectionContextFactory socketConnectionContextFactory, IHostApplicationLifetime lifetime) : IKestrelMiddleware
{
    public async Task InvokeAsync(ConnectionDelegate next, ConnectionContext ctx)
    {
        var result = await ctx.Transport.Input.ReadAsync(lifetime.ApplicationStopping);
        var buffer = result.Buffer;
        var span = buffer.FirstSpan;
        if (!Trojan.TryParseHeader(span, out _, out var command, out var endpoint, out var offset)) {
            ctx.Transport.Input.AdvanceTo(buffer.Start);
            await next(ctx);
            return;
        }
        ctx.Transport.Input.AdvanceTo(buffer.GetPosition(offset));
        if (command == 0x01) {
            await TcpForwardingAsync(ctx, endpoint, lifetime.ApplicationStopping);
        } else {
            await UdpForwardingAsync(ctx, endpoint, lifetime.ApplicationStopping);
        }
    }

    private async Task TcpForwardingAsync(ConnectionContext ctx, EndPoint endpoint, CancellationToken ct)
    {
        using var socket = new Socket(SocketType.Stream, ProtocolType.Tcp);
        socket.NoDelay = true;
        socket.LingerState = new LingerOption(true, 0);
        using (var cts = CancellationTokenSource.CreateLinkedTokenSource(ctx.ConnectionClosed, ct)) {
            cts.CancelAfter(TimeSpan.FromSeconds(6));
            await socket.ConnectAsync(endpoint, cts.Token);
        }
        await using var connection = socketConnectionContextFactory.Create(socket);
        using (var cts = CancellationTokenSource.CreateLinkedTokenSource(ctx.ConnectionClosed, ct)) {
            try {
                await Task.WhenAny(ctx.Transport.Input.CopyToAsync(connection.Transport.Output, cts.Token), connection.Transport.Input.CopyToAsync(ctx.Transport.Output, cts.Token));
            } finally {
                await cts.CancelAsync();
            }
        }
    }

    private static async Task UdpForwardingAsync(ConnectionContext ctx, EndPoint endpoint, CancellationToken ct)
    {
        using var socket = new Socket(SocketType.Dgram, ProtocolType.Udp);
        await socket.ConnectAsync(endpoint, ctx.ConnectionClosed);
        using var cts = CancellationTokenSource.CreateLinkedTokenSource(ctx.ConnectionClosed, ct);
        try {
            await Task.WhenAny(Trojan.Udp.CopyToAsync(ctx, socket, ct), Trojan.Udp.CopyToAsync(socket, ctx, ct));
        } finally {
            await cts.CancelAsync();
        }
    }
}