using System.Collections.Concurrent;
using System.Diagnostics;
using System.Reflection;
using Motus;
using Motus.Abstractions;
using Motus.Cli.Services.Reporters;
using CliTestResult = Motus.Cli.Services.TestResult;

namespace Motus.Cli.Services;

public sealed record TestResult(string FullName, bool Passed, TimeSpan Duration, string? ErrorMessage, string? StackTrace, bool Flaky = false, bool Quarantined = false, int Attempts = 1);

public sealed record TestRunResult(int Total, int Passed, int Failed, int Skipped, TimeSpan Duration, bool CoverageThresholdsFailed = false, int Flaky = 0, int Quarantined = 0);

/// <summary>
/// Controls which failures the runner re-runs under <c>--retries</c>.
/// <see cref="Transient"/> retries only transient CDP disconnects (the default, safe
/// behavior). <see cref="Flake"/> retries any failure and reports a test that fails then
/// passes within its retry budget as flaky rather than failed.
/// </summary>
public enum RetryPolicy
{
    Transient,
    Flake,
}

public sealed class TestRunner(int maxWorkers)
{
    private readonly ConcurrentDictionary<Assembly, bool> _initializedAssemblies = new();
    private readonly ConcurrentDictionary<Type, bool> _initializedClasses = new();

    public Task<TestRunResult> RunAsync(List<DiscoveredTest> tests, IReporter reporter) =>
        RunAsync(tests, reporter, a11yMode: null, enforcePerfBudget: false, coverageReporters: null);

    public Task<TestRunResult> RunAsync(List<DiscoveredTest> tests, IReporter reporter, string? a11yMode) =>
        RunAsync(tests, reporter, a11yMode, enforcePerfBudget: false, coverageReporters: null);

    public Task<TestRunResult> RunAsync(List<DiscoveredTest> tests, IReporter reporter, string? a11yMode, bool enforcePerfBudget) =>
        RunAsync(tests, reporter, a11yMode, enforcePerfBudget, coverageReporters: null);

    public async Task<TestRunResult> RunAsync(
        List<DiscoveredTest> tests,
        IReporter reporter,
        string? a11yMode,
        bool enforcePerfBudget,
        IReadOnlyList<ICoverageReporter>? coverageReporters,
        int maxRetries = 0,
        RetryPolicy retryPolicy = RetryPolicy.Transient,
        int? shardIndex = null,
        int? shardTotal = null)
    {
        if (maxRetries < 0)
            maxRetries = 0;
        var suiteName = tests.Count > 0
            ? tests[0].TestClass.Assembly.GetName().Name ?? "Motus Tests"
            : "Motus Tests";

        await reporter.OnTestRunStartAsync(
            new TestSuiteInfo(suiteName, tests.Count, ShardIndex: shardIndex, ShardTotal: shardTotal));

        var sw = Stopwatch.StartNew();

        // Run [AssemblyInitialize] for each assembly represented in this batch
        var failedAssemblies = new HashSet<Assembly>();
        var assemblies = tests.Select(t => t.TestClass.Assembly).Distinct().ToList();
        foreach (var assembly in assemblies)
        {
            if (_initializedAssemblies.TryAdd(assembly, true))
            {
                try
                {
                    await RunAssemblyInitializeAsync(assembly);
                }
                catch (Exception ex)
                {
                    failedAssemblies.Add(assembly);
                    var msg = ex.InnerException?.Message ?? ex.Message;
                    Console.Error.WriteLine($"AssemblyInitialize failed for {assembly.GetName().Name}: {msg}");
                }
            }
        }

        var semaphore = new SemaphoreSlim(maxWorkers);
        var results = new List<CliTestResult>();
        var collectedCoverage = new List<CoverageData>();
        var lockObj = new object();

        var tasks = tests.Select(async test =>
        {
            // Skip ignored tests entirely — don't fire reporter events
            if (test.IsIgnored)
                return;

            await semaphore.WaitAsync();
            try
            {
                var testInfo = new TestInfo(test.FullName, suiteName);
                await reporter.OnTestStartAsync(testInfo);

                CliTestResult result;
                var attempts = 1;

                if (failedAssemblies.Contains(test.TestClass.Assembly))
                {
                    AccessibilityViolationSink.Begin();
                    PerformanceMetricsSink.Begin();
                    CoverageSink.Begin();
                    result = new CliTestResult(test.FullName, false, TimeSpan.Zero,
                        "AssemblyInitialize failed", null);
                }
                else
                {
                    var totalAttempts = maxRetries + 1;
                    var attempt = 0;
                    Exception? failure;
                    while (true)
                    {
                        attempt++;

                        AccessibilityViolationSink.Begin();
                        PerformanceMetricsSink.Begin();
                        CoverageSink.Begin();

                        // Each call rebuilds the test instance (and its browser context),
                        // so every retry runs against clean state.
                        (result, failure) = await ExecuteTestAsync(test);

                        if (result.Passed || attempt >= totalAttempts || !ShouldRetry(failure, retryPolicy))
                            break;

                        // Discard sink data captured during the failed attempt; a clean
                        // Begin() runs at the top of the next iteration.
                        AccessibilityViolationSink.End();
                        PerformanceMetricsSink.End();
                        CoverageSink.End();

                        WriteRetryNotice(test.FullName, attempt + 1, totalAttempts, retryPolicy);
                    }

                    attempts = attempt;
                }

                var violations = AccessibilityViolationSink.End();
                var perfMetrics = PerformanceMetricsSink.End();
                var coverage = CoverageSink.End();

                if (coverage is not null)
                {
                    lock (lockObj)
                    {
                        collectedCoverage.Add(coverage);
                    }

                    if (coverageReporters is not null)
                    {
                        foreach (var cr in coverageReporters)
                        {
                            try { await cr.OnCoverageCollectedAsync(coverage, testInfo); }
                            catch { }
                        }
                    }

                    if (reporter is ICoverageReporter testReporterCoverage)
                    {
                        try { await testReporterCoverage.OnCoverageCollectedAsync(coverage, testInfo); }
                        catch { }
                    }
                }

                // Dispatch violations to IAccessibilityReporter reporters
                if (violations.Count > 0 && reporter is IAccessibilityReporter a11yReporter)
                {
                    foreach (var violation in violations)
                    {
                        try { await a11yReporter.OnAccessibilityViolationAsync(violation, testInfo); }
                        catch { }
                    }
                }

                // Dispatch performance metrics to IPerformanceReporter reporters
                if (perfMetrics is not null)
                {
                    var budget = ConfigMerge.ToBudget(MotusConfigLoader.Config.Performance);
                    var budgetResult = budget?.Evaluate(perfMetrics);

                    if (reporter is IPerformanceReporter perfReporter)
                    {
                        try { await perfReporter.OnPerformanceMetricsCollectedAsync(perfMetrics, budgetResult, testInfo); }
                        catch { }
                    }

                    // In enforce mode, override a passing test to failed when budget is exceeded
                    if (result.Passed && enforcePerfBudget && budgetResult is { Passed: false })
                    {
                        var failedCount = budgetResult.Entries.Count(e => !e.Passed);
                        result = result with
                        {
                            Passed = false,
                            ErrorMessage = $"Performance budget exceeded: {failedCount} metric(s) over budget.",
                        };
                    }
                }

                // In enforce mode, override a passing test to failed when Error violations exist
                if (result.Passed
                    && string.Equals(a11yMode, "enforce", StringComparison.OrdinalIgnoreCase)
                    && violations.Any(v => v.Severity == AccessibilityViolationSeverity.Error))
                {
                    var count = violations.Count(v => v.Severity == AccessibilityViolationSeverity.Error);
                    result = result with
                    {
                        Passed = false,
                        ErrorMessage = $"Accessibility enforcement failed: {count} error-severity violation(s) detected.",
                    };
                }

                // Classify after any enforce override above: a test the body passed on a
                // later attempt but enforce then failed is a hard failure, not flaky.
                var flaky = result.Passed && attempts > 1;
                result = result with
                {
                    Flaky = flaky,
                    Quarantined = test.Quarantined,
                    Attempts = attempts,
                };

                lock (lockObj)
                {
                    results.Add(result);
                }

                var absResult = new Abstractions.TestResult(
                    result.FullName, result.Passed, result.Duration.TotalMilliseconds,
                    result.ErrorMessage, result.StackTrace,
                    Flaky: result.Flaky, Quarantined: result.Quarantined, Attempts: result.Attempts);
                await reporter.OnTestEndAsync(testInfo, absResult);
            }
            finally
            {
                semaphore.Release();
            }
        });

        await Task.WhenAll(tasks);
        sw.Stop();

        // Aggregate coverage and dispatch run-end events
        CoverageData? aggregated = null;
        var coverageThresholdsFailed = false;
        if (collectedCoverage.Count > 0)
        {
            var mergedScripts = CoverageAggregator.MergeScripts(collectedCoverage.SelectMany(c => c.Scripts));
            var mergedSheets = CoverageAggregator.MergeStylesheets(collectedCoverage.SelectMany(c => c.Stylesheets));
            var mergedOriginal = CoverageAggregator.MergeOriginalFiles(collectedCoverage.SelectMany(c => c.OriginalFiles));
            var summary = CoverageAggregator.BuildSummary(mergedScripts, mergedSheets);
            var diagnostic = collectedCoverage.Select(c => c.DiagnosticMessage).FirstOrDefault(d => !string.IsNullOrEmpty(d));

            aggregated = new CoverageData(
                mergedScripts,
                mergedSheets,
                summary,
                DateTime.UtcNow,
                diagnostic,
                mergedOriginal);

            // Threshold enforcement: read from MotusConfig and force-fail when violated
            var coverageOptions = ConfigMerge.ToCoverageOptions(MotusConfigLoader.Config.Coverage);
            if (coverageOptions is not null)
            {
                var thresholdResult = CoverageThresholds.Evaluate(aggregated, coverageOptions);
                if (!thresholdResult.Passed)
                {
                    coverageThresholdsFailed = true;
                    foreach (var entry in thresholdResult.Failed)
                    {
                        Console.Error.WriteLine(
                            $"Coverage threshold failed: {entry.MetricName} {entry.Actual:F1}% < {entry.Threshold:F1}%");
                    }
                }
            }

            if (coverageReporters is not null)
            {
                foreach (var cr in coverageReporters)
                {
                    try { await cr.OnCoverageRunEndAsync(aggregated); }
                    catch { }
                }
            }

            if (reporter is ICoverageReporter testReporterCoverageEnd)
            {
                try { await testReporterCoverageEnd.OnCoverageRunEndAsync(aggregated); }
                catch { }
            }
        }

        // Run [ClassCleanup] for all initialized classes
        foreach (var type in _initializedClasses.Keys)
        {
            try
            {
                await RunStaticLifecycleMethodAsync(type, "ClassCleanupAttribute");
                await RunStaticLifecycleMethodAsync(type, "OneTimeTearDownAttribute");
            }
            catch (Exception ex)
            {
                Console.Error.WriteLine($"ClassCleanup failed for {type.FullName}: {ex.InnerException?.Message ?? ex.Message}");
            }
        }

        // Run [AssemblyCleanup] for each initialized assembly
        foreach (var assembly in assemblies)
        {
            try
            {
                await RunAssemblyCleanupAsync(assembly);
            }
            catch (Exception ex)
            {
                Console.Error.WriteLine($"AssemblyCleanup failed for {assembly.GetName().Name}: {ex.InnerException?.Message ?? ex.Message}");
            }
        }

        var skipped = tests.Count(t => t.IsIgnored);
        // Quarantined results occupy their own bucket and are excluded from passed/failed
        // so a quarantined failure never gates the run.
        var quarantined = results.Count(r => r.Quarantined);
        var passed = results.Count(r => r.Passed && !r.Quarantined);
        var failed = results.Count(r => !r.Passed && !r.Quarantined);
        var flaky = results.Count(r => r.Flaky && !r.Quarantined);
        var runResult = new TestRunResult(
            passed + failed + skipped + quarantined, passed, failed, skipped, sw.Elapsed,
            coverageThresholdsFailed, flaky, quarantined);

        await reporter.OnTestRunEndAsync(new TestRunSummary(
            suiteName, passed, failed, skipped, sw.Elapsed.TotalMilliseconds, flaky, quarantined));
        return runResult;
    }

    /// <summary>
    /// Emits a [RETRY] notice in the same shape (and color, when stderr is a TTY)
    /// as the [PASS]/[FAIL] lines emitted by ConsoleReporter, so the retry log
    /// reads as part of the test stream instead of free-form debug output.
    /// </summary>
    private static void WriteRetryNotice(string testName, int attempt, int totalAttempts, RetryPolicy policy)
    {
        const string Yellow = "\x1b[33m";
        const string Gray = "\x1b[90m";
        const string Reset = "\x1b[0m";

        var reason = policy == RetryPolicy.Flake ? "failure" : "CDP disconnect";

        var useColor = !Console.IsErrorRedirected;
        if (useColor)
        {
            Console.Error.WriteLine(
                $"  {Yellow}[RETRY]{Reset} {testName} {Gray}({reason}, attempt {attempt}/{totalAttempts}){Reset}");
        }
        else
        {
            Console.Error.WriteLine(
                $"  [RETRY] {testName} ({reason}, attempt {attempt}/{totalAttempts})");
        }
    }

    /// <summary>
    /// Decides whether a failed attempt should be retried under the active policy.
    /// <see cref="RetryPolicy.Flake"/> retries any failure; <see cref="RetryPolicy.Transient"/>
    /// retries only transient CDP disconnects.
    /// </summary>
    private static bool ShouldRetry(Exception? failure, RetryPolicy policy)
    {
        if (failure is null)
            return false;
        return policy == RetryPolicy.Flake || IsTransientCdpFailure(failure);
    }

    /// <summary>
    /// Returns true when the exception chain contains a transient CDP failure that
    /// is safe to recover from by re-running the entire test (which rebuilds a fresh
    /// browser context and WebSocket). Used to drive --retries.
    /// </summary>
    private static bool IsTransientCdpFailure(Exception? ex)
    {
        while (ex is not null)
        {
            if (ex is CdpDisconnectedException || ex is Abstractions.MotusTargetClosedException)
                return true;
            ex = ex.InnerException;
        }
        return false;
    }

    private async Task<(CliTestResult Result, Exception? Failure)> ExecuteTestAsync(DiscoveredTest test)
    {
        var testSw = Stopwatch.StartNew();
        object? instance = null;
        try
        {
            instance = Activator.CreateInstance(test.TestClass)!;

            // Run [ClassInitialize] once per type (static, takes TestContext)
            if (_initializedClasses.TryAdd(test.TestClass, true))
            {
                await RunStaticLifecycleMethodAsync(test.TestClass, "ClassInitializeAttribute");
                await RunStaticLifecycleMethodAsync(test.TestClass, "OneTimeSetUpAttribute");
            }

            // Run [TestInitialize] / [SetUp]
            await ExecuteLifecycleMethodAsync(instance, test.TestClass, "TestInitializeAttribute");
            await ExecuteLifecycleMethodAsync(instance, test.TestClass, "SetUpAttribute");

            var result = test.TestMethod.Invoke(instance, null);

            if (result is Task task)
                await task;

            testSw.Stop();
            return (new CliTestResult(test.FullName, true, testSw.Elapsed, null, null), null);
        }
        catch (TargetInvocationException ex) when (ex.InnerException is not null)
        {
            testSw.Stop();
            return (
                new CliTestResult(test.FullName, false, testSw.Elapsed, ex.InnerException.Message, ex.InnerException.StackTrace),
                ex.InnerException);
        }
        catch (Exception ex)
        {
            testSw.Stop();
            return (
                new CliTestResult(test.FullName, false, testSw.Elapsed, ex.Message, ex.StackTrace),
                ex);
        }
        finally
        {
            if (instance is not null)
            {
                try
                {
                    // Run [TestCleanup] / [TearDown]
                    await ExecuteLifecycleMethodAsync(instance, test.TestClass, "TestCleanupAttribute");
                    await ExecuteLifecycleMethodAsync(instance, test.TestClass, "TearDownAttribute");
                }
                catch (Exception ex)
                {
                    Console.Error.WriteLine($"TestCleanup failed for {test.FullName}: {ex.InnerException?.Message ?? ex.Message}");
                }

                if (instance is IAsyncDisposable asyncDisposable)
                    await asyncDisposable.DisposeAsync();
                else if (instance is IDisposable disposable)
                    disposable.Dispose();
            }
        }
    }

    private static async Task ExecuteLifecycleMethodAsync(object instance, Type testClass, string attributeName)
    {
        // Walk the hierarchy bottom-up, then reverse so base-class methods run first
        var methods = new List<MethodInfo>();
        var type = testClass;
        while (type is not null && type != typeof(object))
        {
            var declared = type.GetMethods(BindingFlags.Public | BindingFlags.Instance | BindingFlags.DeclaredOnly)
                .Where(m => m.GetCustomAttributes(true).Any(a => a.GetType().Name == attributeName));
            methods.InsertRange(0, declared);
            type = type.BaseType;
        }

        foreach (var method in methods)
        {
            var result = method.Invoke(instance, null);
            if (result is Task task)
                await task;
        }
    }

    private static async Task RunAssemblyCleanupAsync(Assembly assembly)
    {
        foreach (var type in assembly.GetExportedTypes())
        {
            var methods = type.GetMethods(BindingFlags.Public | BindingFlags.Static)
                .Where(m => m.GetCustomAttributes(true).Any(a => a.GetType().Name == "AssemblyCleanupAttribute"));

            foreach (var method in methods)
            {
                var result = method.Invoke(null, null);
                if (result is Task task)
                    await task;
            }
        }
    }

    private static async Task RunAssemblyInitializeAsync(Assembly assembly)
    {
        foreach (var type in assembly.GetExportedTypes())
        {
            var methods = type.GetMethods(BindingFlags.Public | BindingFlags.Static)
                .Where(m => m.GetCustomAttributes(true).Any(a => a.GetType().Name == "AssemblyInitializeAttribute"));

            foreach (var method in methods)
            {
                var parameters = method.GetParameters();
                object?[] args = parameters.Length > 0 ? [null] : [];
                var result = method.Invoke(null, args);
                if (result is Task task)
                    await task;
            }
        }
    }

    private static async Task RunStaticLifecycleMethodAsync(Type type, string attributeName)
    {
        var methods = type.GetMethods(BindingFlags.Public | BindingFlags.Static | BindingFlags.FlattenHierarchy)
            .Where(m => m.GetCustomAttributes(true).Any(a => a.GetType().Name == attributeName));

        foreach (var method in methods)
        {
            var parameters = method.GetParameters();
            object?[] args = parameters.Length > 0 ? [null] : [];
            var result = method.Invoke(null, args);
            if (result is Task task)
                await task;
        }
    }
}
