using System.Net;
using System.Net.Sockets;
using System.Text;

namespace Agent.Core.Tests;

internal sealed class SharedFakeOllamaServer : IAsyncDisposable
{
    private readonly TcpListener _listener = new(IPAddress.Loopback, 0);
    private readonly CancellationTokenSource _stop = new();
    private readonly Task _loop;

    public SharedFakeOllamaServer()
    {
        _listener.Start();
        Endpoint = new Uri($"http://127.0.0.1:{((IPEndPoint)_listener.LocalEndpoint).Port}/");
        _loop = RunAsync();
    }

    public Uri Endpoint { get; }

    private async Task RunAsync()
    {
        try
        {
            while (!_stop.IsCancellationRequested)
            {
                using TcpClient client = await _listener.AcceptTcpClientAsync(_stop.Token).ConfigureAwait(false);
                NetworkStream stream = client.GetStream();
                using var reader = new StreamReader(stream, Encoding.ASCII, false, 1024, true);
                int contentLength = 0;
                while (await reader.ReadLineAsync(_stop.Token).ConfigureAwait(false) is { } line && line.Length > 0)
                    if (line.StartsWith("Content-Length:", StringComparison.OrdinalIgnoreCase)) contentLength = int.Parse(line[15..].Trim(), System.Globalization.CultureInfo.InvariantCulture);
                if (contentLength > 0)
                {
                    char[] body = new char[contentLength];
                    await reader.ReadBlockAsync(body.AsMemory(), _stop.Token).ConfigureAwait(false);
                }
                byte[] payload = Encoding.UTF8.GetBytes("""{"message":{"role":"assistant","content":"done"},"prompt_eval_count":3,"eval_count":2}""");
                byte[] header = Encoding.ASCII.GetBytes($"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {payload.Length}\r\nConnection: close\r\n\r\n");
                await stream.WriteAsync(header, _stop.Token).ConfigureAwait(false);
                await stream.WriteAsync(payload, _stop.Token).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException) when (_stop.IsCancellationRequested) { }
    }

    public async ValueTask DisposeAsync()
    {
        await _stop.CancelAsync().ConfigureAwait(false);
        _listener.Stop();
        await _loop.ConfigureAwait(false);
        _stop.Dispose();
    }
}