using System.Collections.Concurrent;
using System.Text.Json;
using System.Threading.Channels;

namespace Motus;

/// <summary>
/// Raw event payload surfaced to event channels before typed deserialization.
/// </summary>
internal readonly record struct RawCdpEvent(JsonElement Params, string? SessionId);

/// <summary>
/// Core CDP WebSocket transport. Manages a single WebSocket connection, a background
/// receive loop, request/response correlation, and event channel dispatch.
/// </summary>
internal sealed class CdpTransport : IMotusTransport
{
    private readonly ICdpSocket _socket;
    private readonly TimeSpan _slowMo;
    private readonly SemaphoreSlim _sendLock = new(1, 1);
    private readonly CancellationTokenSource _cts = new();
    private Task? _receiveLoop;
    private bool _disposed;
    // Set as soon as the receive loop observes a disconnect (clean close, exception,
    // or socket teardown). Distinct from _disposed: a disconnect can be detected
    // while the transport is still owned by callers that haven't disposed it yet.
    // Without this flag, SendRawAsync would happily enqueue new commands after the
    // receive loop has exited, then wait the full CommandTimeout (60s) per command
    // for responses that can never arrive — turning post-disconnect cleanup paths
    // (e.g. coverage teardown firing many getScriptSource calls) into multi-minute
    // hangs.
    private volatile bool _disconnected;

    /// <summary>
    /// Maximum time to wait for a CDP response before treating the browser as unresponsive.
    /// Prevents indefinite hangs when Chrome freezes without closing the WebSocket.
    /// </summary>
    private static readonly TimeSpan CommandTimeout = TimeSpan.FromSeconds(60);

    private readonly ConcurrentDictionary<int, TaskCompletionSource<JsonElement>> _pending = new();
    private readonly ConcurrentDictionary<string, Channel<RawCdpEvent>> _eventChannels = new();
    private int _nextId;

    /// <inheritdoc />
    public MotusCapabilities Capabilities => MotusCapabilities.AllCdp;

    /// <summary>
    /// Raised when the WebSocket connection is lost. The exception (if any) is provided.
    /// </summary>
    public event Action<Exception?>? Disconnected;

    internal CdpTransport(ICdpSocket socket, TimeSpan slowMo = default)
    {
        _socket = socket;
        _slowMo = slowMo;
    }

    /// <summary>
    /// Connects to the browser CDP endpoint and starts the background receive loop.
    /// </summary>
    internal async Task ConnectAsync(Uri endpointUri, CancellationToken ct)
    {
        await _socket.ConnectAsync(endpointUri, ct).ConfigureAwait(false);
        _receiveLoop = RunReceiveLoopAsync(_cts.Token);
    }

    /// <summary>
    /// Sends a raw CDP command and awaits the correlated response.
    /// </summary>
    internal async Task<JsonElement> SendRawAsync(
        string method, JsonElement paramsElement, string? sessionId, CancellationToken ct)
    {
        if (_disposed || _disconnected)
            throw new CdpDisconnectedException();

        var id = Interlocked.Increment(ref _nextId);
        var tcs = new TaskCompletionSource<JsonElement>(TaskCreationOptions.RunContinuationsAsynchronously);
        _pending[id] = tcs;

        // Ensure commands always have a timeout, even if the caller passes CancellationToken.None.
        // This prevents indefinite hangs when Chrome becomes unresponsive without disconnecting.
        using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        timeoutCts.CancelAfter(CommandTimeout);
        var effectiveCt = timeoutCts.Token;

        try
        {
            if (_slowMo > TimeSpan.Zero)
                await Task.Delay(_slowMo, effectiveCt).ConfigureAwait(false);

            var envelope = new CdpCommandEnvelope(id, method, paramsElement, sessionId);
            var bytes = JsonSerializer.SerializeToUtf8Bytes(envelope, CdpJsonContext.Default.CdpCommandEnvelope);

            // ClientWebSocket does not support concurrent sends; serialize access.
            await _sendLock.WaitAsync(effectiveCt).ConfigureAwait(false);
            try
            {
                await _socket.SendAsync(bytes, effectiveCt).ConfigureAwait(false);
            }
            finally
            {
                _sendLock.Release();
            }

            // Link cancellation (caller + timeout) to the pending TCS
            await using var reg = effectiveCt.Register(() => tcs.TrySetCanceled(effectiveCt));
            return await tcs.Task.ConfigureAwait(false);
        }
        finally
        {
            _pending.TryRemove(id, out _);
        }
    }

    /// <summary>
    /// Gets or creates an event channel for the given composite key.
    /// The key format is <c>"Domain.eventName|sessionId"</c> (empty string for browser-level).
    /// </summary>
    internal Channel<RawCdpEvent> GetOrCreateEventChannel(string channelKey)
    {
        return _eventChannels.GetOrAdd(channelKey, static _ =>
            Channel.CreateUnbounded<RawCdpEvent>(new UnboundedChannelOptions
            {
                SingleReader = false,
                SingleWriter = true
            }));
    }

    /// <summary>
    /// Removes and completes all event channels associated with the given session ID.
    /// Called when a page/session is disposed to prevent unbounded channel accumulation.
    /// </summary>
    internal void RemoveChannelsForSession(string sessionId)
    {
        var suffix = $"|{sessionId}";
        foreach (var key in _eventChannels.Keys)
        {
            if (key.EndsWith(suffix, StringComparison.Ordinal) &&
                _eventChannels.TryRemove(key, out var channel))
            {
                channel.Writer.TryComplete();
            }
        }
    }

    private async Task RunReceiveLoopAsync(CancellationToken ct)
    {
        try
        {
            while (!ct.IsCancellationRequested)
            {
                var message = await _socket.ReceiveAsync(ct).ConfigureAwait(false);

                if (message.IsEmpty)
                {
                    // Clean disconnect
                    MarkDisconnectedAndFault(new CdpDisconnectedException());
                    Disconnected?.Invoke(null);
                    return;
                }

                // Copy from the shared socket buffer so the next ReceiveAsync
                // cannot overwrite data still referenced by JsonElement values.
                var messageBytes = message.ToArray();

                try
                {
                    var envelope = JsonSerializer.Deserialize(
                        messageBytes.AsSpan(),
                        CdpJsonContext.Default.CdpInboundEnvelope);

                    if (envelope is not null)
                        DispatchMessage(envelope);
                }
                catch (JsonException)
                {
                    // Malformed CDP message; skip it.
                }
            }
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            // Normal shutdown via dispose
        }
        catch (Exception ex)
        {
            // Unexpected disconnect
            MarkDisconnectedAndFault(new CdpDisconnectedException(ex));
            Disconnected?.Invoke(ex);
        }
    }

    private void DispatchMessage(CdpInboundEnvelope envelope)
    {
        if (envelope.Id is not null)
        {
            DispatchResponse(envelope);
        }
        else if (envelope.Method is not null)
        {
            DispatchEvent(envelope);
        }
    }

    private void DispatchResponse(CdpInboundEnvelope envelope)
    {
        if (!_pending.TryRemove(envelope.Id!.Value, out var tcs))
            return;

        if (envelope.Error is not null)
        {
            tcs.TrySetException(new CdpProtocolException(envelope.Error.Code, envelope.Error.Message));
        }
        else
        {
            // Default to an empty JSON object if result is missing
            tcs.TrySetResult(envelope.Result ?? EmptyJsonElement());
        }
    }

    private void DispatchEvent(CdpInboundEnvelope envelope)
    {
        var sessionId = envelope.SessionId ?? string.Empty;
        var channelKey = $"{envelope.Method}|{sessionId}";

        if (_eventChannels.TryGetValue(channelKey, out var channel))
        {
            var rawEvent = new RawCdpEvent(
                envelope.Params ?? default,
                envelope.SessionId);

            channel.Writer.TryWrite(rawEvent);
        }
    }

    private static readonly JsonElement s_emptyObject = JsonDocument.Parse("{}").RootElement.Clone();

    internal static JsonElement EmptyJsonElement() => s_emptyObject;

    private void FaultAllPending(Exception exception)
    {
        foreach (var kvp in _pending)
        {
            if (_pending.TryRemove(kvp.Key, out var tcs))
                tcs.TrySetException(exception);
        }
    }

    /// <summary>
    /// Single entry point for "the connection is gone": flips <see cref="_disconnected"/>
    /// so subsequent <see cref="SendRawAsync"/> calls fail fast, faults pending requests
    /// with the supplied exception, and completes all event channels.
    /// </summary>
    private void MarkDisconnectedAndFault(Exception exception)
    {
        _disconnected = true;
        FaultAllPending(exception);
        CompleteAllChannels();
    }

    private void CompleteAllChannels()
    {
        foreach (var kvp in _eventChannels)
        {
            kvp.Value.Writer.TryComplete();
        }
    }

    public async ValueTask DisposeAsync()
    {
        if (_disposed) return;
        _disposed = true;

        await _cts.CancelAsync().ConfigureAwait(false);

        if (_receiveLoop is not null)
        {
            try { await _receiveLoop.ConfigureAwait(false); }
            catch { /* Swallow; loop already handled cleanup */ }
        }

        MarkDisconnectedAndFault(new CdpDisconnectedException());

        await _socket.DisposeAsync().ConfigureAwait(false);
        _cts.Dispose();
        _sendLock.Dispose();
    }
}
