using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using System.Threading.Channels;

namespace Backend;

/// <summary>
/// One <see cref="GameSession"/> per connected browser. Runs the physics tick loop
/// and the WebSocket reader loop concurrently until either side disconnects.
/// </summary>
public sealed class GameSession : IAsyncDisposable
{
    private readonly WebSocket _socket;
    private readonly GameProcess _game;
    private readonly PolicyRunner _policy;
    private readonly ILogger<GameSession> _logger;

    // Mailbox: cursor updates from the browser. Bounded to 1 with drop-newest
    // isn't quite right — we want drop-oldest so the ticker always reads the
    // latest cursor. Channels' DropWrite behavior keeps the first, so we do
    // it manually by draining the reader on each tick.
    private readonly Channel<(float X, float Y)> _cursorInbox =
        Channel.CreateUnbounded<(float, float)>();

    private (float X, float Y) _currentCursor = (0.5f, 0.5f);
    private float[] _lastObs = new float[PolicyRunner.ObsDim];

    public GameSession(WebSocket socket, GameProcess game, PolicyRunner policy,
                      ILogger<GameSession> logger)
    {
        _socket = socket;
        _game = game;
        _policy = policy;
        _logger = logger;
    }

    public async Task RunAsync(CancellationToken cancellation)
    {
        await _game.StartAsync();

        // First step (noop) to get an initial observation for the policy.
        var initial = await _game.ResetAsync();
        _lastObs = initial.Observation;

        // Send init frame to the browser.
        await SendJsonAsync(new InitFrame { World = new[] { 1f, 1f } },
            WireJsonContext.Default.InitFrame, cancellation);

        var readerTask = ReaderLoopAsync(cancellation);
        var tickerTask = TickerLoopAsync(cancellation);

        // First loop to complete cancels the other.
        var done = await Task.WhenAny(readerTask, tickerTask);
        try { await done; }
        catch (OperationCanceledException) { /* expected */ }
        catch (Exception ex) { _logger.LogWarning(ex, "session loop ended"); }
    }

    private async Task ReaderLoopAsync(CancellationToken cancellation)
    {
        var buffer = new byte[4 * 1024];
        while (!cancellation.IsCancellationRequested)
        {
            var result = await _socket.ReceiveAsync(buffer, cancellation);
            if (result.MessageType == WebSocketMessageType.Close) break;
            if (result.MessageType != WebSocketMessageType.Text) continue;

            var text = Encoding.UTF8.GetString(buffer, 0, result.Count);
            CursorMessage? msg;
            try
            {
                msg = JsonSerializer.Deserialize(text, WireJsonContext.Default.CursorMessage);
            }
            catch (JsonException) { continue; }
            if (msg is null || msg.Type != "cursor") continue;

            await _cursorInbox.Writer.WriteAsync(
                (Clamp01(msg.X), Clamp01(msg.Y)), cancellation);
        }
    }

    private async Task TickerLoopAsync(CancellationToken cancellation)
    {
        using var timer = new PeriodicTimer(TimeSpan.FromMilliseconds(20));  // 50 Hz
        while (await timer.WaitForNextTickAsync(cancellation))
        {
            // Drain the cursor mailbox — keep only the latest update.
            while (_cursorInbox.Reader.TryRead(out var next)) _currentCursor = next;

            int action = _policy.SelectAction(_lastObs);
            var step = await _game.StepAsync(action, _currentCursor);
            _lastObs = step.Observation;

            var frame = new StateFrame
            {
                X = step.State.X,
                Y = step.State.Y,
                Angle = step.State.Angle,
                Engine = step.State.Engine,
                Target = new[] { _currentCursor.X, _currentCursor.Y },
                Step = step.Step,
            };
            await SendJsonAsync(frame, WireJsonContext.Default.StateFrame, cancellation);
        }
    }

    private async Task SendJsonAsync<T>(
        T value,
        System.Text.Json.Serialization.Metadata.JsonTypeInfo<T> ctx,
        CancellationToken cancellation)
    {
        var bytes = JsonSerializer.SerializeToUtf8Bytes(value, ctx);
        await _socket.SendAsync(bytes, WebSocketMessageType.Text,
            endOfMessage: true, cancellation);
    }

    private static float Clamp01(float v) => v < 0f ? 0f : (v > 1f ? 1f : v);

    public async ValueTask DisposeAsync()
    {
        try
        {
            if (_socket.State == WebSocketState.Open)
                await _socket.CloseAsync(WebSocketCloseStatus.NormalClosure,
                    "session-end", CancellationToken.None);
        }
        catch { }
        await _game.DisposeAsync();
    }
}