using System.Runtime.CompilerServices;
using System.Text.Json;
using System.Threading.Channels;
using Motus.Abstractions;

namespace Motus;

/// <summary>
/// Real ITracing implementation backed by CDP Tracing domain.
/// Uses the browser-level CDP session since Tracing is a browser-target domain.
/// </summary>
internal sealed class Tracing : ITracing
{
    // CDP Tracing is browser-wide: only one trace session can be active per browser.
    // When multiple BrowserContexts share a browser (e.g. shared-browser test fixtures
    // with parallel workers), concurrent StartAsync calls collide. Serialize per session.
    private static readonly ConditionalWeakTable<IMotusSession, SemaphoreSlim> s_browserGates = new();

    // Upper bound on how long StartAsync will wait for the browser-wide trace gate
    // before failing. Exists to convert a leaked gate (test exited mid-trace without
    // reaching StopAsync, leaving the semaphore at count 0 forever) into an actionable
    // error rather than an indefinite test-run hang. A real trace shouldn't take
    // anywhere near this; longer-running scenarios should still serialize cleanly.
    private static readonly TimeSpan BrowserGateTimeout = TimeSpan.FromMinutes(2);

    private readonly IMotusSession _browserSession;
    private readonly SemaphoreSlim _browserGate;
    private readonly Channel<JsonElement[]> _dataChannel = Channel.CreateUnbounded<JsonElement[]>();
    private TaskCompletionSource<TracingTracingCompleteEvent>? _completeTcs;
    private CancellationTokenSource? _pumpCts;
    private int _started;
    private bool _holdsBrowserGate;

    internal Tracing(IMotusSession browserSession)
    {
        _browserSession = browserSession;
        _browserGate = s_browserGates.GetValue(browserSession, _ => new SemaphoreSlim(1, 1));
    }

    /// <summary>
    /// Optional HarLog to include in the trace ZIP when stopping.
    /// Set by BrowserContext when HAR recording is active.
    /// </summary>
    internal HarLog? HarLog { get; set; }

    public async Task StartAsync(TracingStartOptions? options = null)
    {
        if (Interlocked.CompareExchange(ref _started, 1, 0) != 0)
            throw new InvalidOperationException("Tracing is already started.");

        CapabilityGuard.Require(_browserSession.Capabilities, MotusCapabilities.Tracing,
            "Tracing", CapabilityGuard.GetTransportDescription(_browserSession));

        // Wait for any other tracing session on the same browser to finish.
        // Bounded so a leaked gate from a previous test that exited mid-trace
        // doesn't deadlock the whole run.
        if (!await _browserGate.WaitAsync(BrowserGateTimeout).ConfigureAwait(false))
        {
            _started = 0;
            throw new TimeoutException(
                $"Timed out after {BrowserGateTimeout.TotalSeconds:F0}s waiting for the browser-wide tracing gate. " +
                "A previous Tracing.StartAsync on this browser likely did not reach StopAsync; " +
                "wrap test bodies that call StartAsync in try/finally to ensure StopAsync runs.");
        }
        _holdsBrowserGate = true;

        options ??= new TracingStartOptions();

        var categories = new List<string> { "devtools.timeline", "-*" };

        if (options.Screenshots == true)
            categories.Add("disabled-by-default-devtools.screenshot");

        if (options.Snapshots == true)
        {
            categories.Add("disabled-by-default-devtools.timeline.layers");
            categories.Add("disabled-by-default-devtools.timeline.picture");
        }

        // Cancel any lingering pumps from a previous run before starting new ones
        _pumpCts?.Cancel();

        // Drain any leftover data from a previous run
        while (_dataChannel.Reader.TryRead(out _)) { }

        var completeTcs = new TaskCompletionSource<TracingTracingCompleteEvent>(
            TaskCreationOptions.RunContinuationsAsynchronously);
        _completeTcs = completeTcs;

        _pumpCts = new CancellationTokenSource();

        // Start background pump for dataCollected events
        _ = PumpDataCollectedAsync(_pumpCts.Token);

        // Subscribe to tracingComplete; capture the local TCS so cancellation
        // of an old pump cannot poison a newer TCS via the shared field.
        _ = PumpTracingCompleteAsync(completeTcs, _pumpCts.Token);

        try
        {
            await _browserSession.SendAsync(
                "Tracing.start",
                new TracingStartParams(
                    TransferMode: "ReturnAsStream",
                    TraceConfig: new TracingTraceConfig(
                        IncludedCategories: categories.ToArray())),
                CdpJsonContext.Default.TracingStartParams,
                CdpJsonContext.Default.TracingStartResult,
                CancellationToken.None).ConfigureAwait(false);
        }
        catch
        {
            // CDP rejected start (e.g. browser-level trace already active outside our control).
            // Roll back state so a retry — or another Tracing instance — can proceed.
            _started = 0;
            _pumpCts?.Cancel();
            _completeTcs?.TrySetCanceled();
            ReleaseBrowserGate();
            throw;
        }
    }

    private void ReleaseBrowserGate()
    {
        if (_holdsBrowserGate)
        {
            _holdsBrowserGate = false;
            _browserGate.Release();
        }
    }

    /// <summary>
    /// Defensive cleanup invoked when the owning <see cref="BrowserContext"/> is being disposed.
    /// If a test called <see cref="StartAsync"/> but exited before reaching <see cref="StopAsync"/>
    /// (early return, untracked exception path, etc.), the browser-wide gate is still held by
    /// this instance. Release it so subsequent tests on the same browser can start tracing
    /// rather than deadlocking on <see cref="BrowserGateTimeout"/>. Best-effort and idempotent.
    /// </summary>
    internal void ReleaseStateOnContextClose()
    {
        try { _pumpCts?.Cancel(); } catch { }
        _completeTcs?.TrySetCanceled();
        _started = 0;
        ReleaseBrowserGate();
    }

    public async Task StopAsync(TracingStopOptions? options = null)
    {
        if (Interlocked.CompareExchange(ref _started, 0, 1) != 1)
            return;

        var completeTcs = _completeTcs!;
        var pumpCts = _pumpCts;

        TracingTracingCompleteEvent completeEvent;
        try
        {
            // Send Tracing.end
            await _browserSession.SendAsync(
                "Tracing.end",
                CdpJsonContext.Default.TracingEndResult,
                CancellationToken.None).ConfigureAwait(false);

            // Wait for tracingComplete event with a timeout to prevent indefinite hangs
            using var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
            using var reg = timeoutCts.Token.Register(() => completeTcs.TrySetCanceled());

            try
            {
                completeEvent = await completeTcs.Task.ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested)
            {
                throw new TimeoutException(
                    "Timed out waiting for Tracing.tracingComplete event after 30 seconds.");
            }
        }
        finally
        {
            // Cancel background pumps and release the browser gate so other
            // contexts can begin tracing even when the stop path errors out.
            pumpCts?.Cancel();
            ReleaseBrowserGate();
        }

        // Collect all accumulated data chunks
        var allEvents = new List<JsonElement>();
        _dataChannel.Writer.TryComplete();
        await foreach (var chunk in _dataChannel.Reader.ReadAllAsync().ConfigureAwait(false))
        {
            allEvents.AddRange(chunk);
        }

        // If a stream handle was returned, read the stream data
        if (completeEvent.Stream is not null)
        {
            var streamEvents = await ReadStreamAsync(completeEvent.Stream).ConfigureAwait(false);
            allEvents.AddRange(streamEvents);
        }

        // Extract screenshots from trace events
        var screenshots = ExtractScreenshots(allEvents);

        // Write trace ZIP if path specified
        if (options?.Path is not null)
        {
            await TracePackager.WriteAsync(
                options.Path,
                allEvents,
                HarLog,
                screenshots).ConfigureAwait(false);
        }

        // Reset for potential reuse
        HarLog = null;
    }

    private async Task PumpDataCollectedAsync(CancellationToken ct)
    {
        try
        {
            await foreach (var evt in _browserSession.SubscribeAsync(
                "Tracing.dataCollected",
                CdpJsonContext.Default.TracingDataCollectedEvent,
                ct).ConfigureAwait(false))
            {
                await _dataChannel.Writer.WriteAsync(evt.Value, ct).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException) { }
    }

    private async Task PumpTracingCompleteAsync(
        TaskCompletionSource<TracingTracingCompleteEvent> completeTcs,
        CancellationToken ct)
    {
        try
        {
            await foreach (var evt in _browserSession.SubscribeAsync(
                "Tracing.tracingComplete",
                CdpJsonContext.Default.TracingTracingCompleteEvent,
                ct).ConfigureAwait(false))
            {
                completeTcs.TrySetResult(evt);
                break;
            }
        }
        catch (OperationCanceledException)
        {
            completeTcs.TrySetCanceled();
        }
    }

    private async Task<List<JsonElement>> ReadStreamAsync(string streamHandle)
    {
        var events = new List<JsonElement>();

        try
        {
            while (true)
            {
                var readResult = await _browserSession.SendAsync(
                    "IO.read",
                    new IoReadParams(streamHandle),
                    CdpJsonContext.Default.IoReadParams,
                    CdpJsonContext.Default.IoReadResult,
                    CancellationToken.None).ConfigureAwait(false);

                if (!string.IsNullOrEmpty(readResult.Data))
                {
                    byte[] bytes;
                    if (readResult.Base64Encoded)
                        bytes = Convert.FromBase64String(readResult.Data);
                    else
                        bytes = System.Text.Encoding.UTF8.GetBytes(readResult.Data);

                    try
                    {
                        using var doc = JsonDocument.Parse(bytes);
                        if (doc.RootElement.ValueKind == JsonValueKind.Array)
                        {
                            foreach (var el in doc.RootElement.EnumerateArray())
                                events.Add(el.Clone());
                        }
                    }
                    catch (JsonException)
                    {
                        // Partial JSON chunk, store as raw string element
                    }
                }

                if (readResult.Eof)
                    break;
            }

            // Close the stream
            await _browserSession.SendAsync(
                "IO.close",
                new IoCloseParams(streamHandle),
                CdpJsonContext.Default.IoCloseParams,
                CdpJsonContext.Default.IoCloseResult,
                CancellationToken.None).ConfigureAwait(false);
        }
        catch (Exception)
        {
            // Stream may already be closed
        }

        return events;
    }

    private static List<ScreenshotEntry> ExtractScreenshots(List<JsonElement> events)
    {
        var screenshots = new List<ScreenshotEntry>();
        int seq = 0;

        foreach (var evt in events)
        {
            if (evt.ValueKind != JsonValueKind.Object)
                continue;

            if (evt.TryGetProperty("cat", out var cat)
                && cat.GetString() == "disabled-by-default-devtools.screenshot"
                && evt.TryGetProperty("args", out var args)
                && args.TryGetProperty("snapshot", out var snapshot))
            {
                var base64 = snapshot.GetString();
                if (base64 is not null)
                {
                    screenshots.Add(new ScreenshotEntry(seq++, Convert.FromBase64String(base64)));
                }
            }
        }

        return screenshots;
    }
}
