using System.Text.Json;
using Microsoft.EntityFrameworkCore;
using w4c_workflows.Data;
using w4c_workflows.Models;
using w4c_workflows.Services.Execution;
using w4c_workflows.Services.Messaging;

namespace w4c_workflows.Services.Runs;

/// <summary>
/// The run lifecycle state machine, extracted from the hosted service so it is
/// unit-testable without a running host. Owns the two async halves of a run:
///
///   <list type="number">
///   <item><b>Dispatch</b> — picks up <c>pending</c> runs, marks them
///   <c>running</c>, and enqueues the start task (root, or <c>StartTaskId</c>
///   for a durable resume) as a <c>task.run</c> job.</item>
///   <item><b>Advance</b> — consumes <c>task.result</c> events from
///   <c>wf:{tenant}:results</c>, updates the TaskRun, and follows the success
///   <c>nextId</c> chain until the run is <c>succeeded</c> or <c>failed</c>.</item>
///   </list>
///
/// On a task failure the saga policy (step 10) applies: retry with exponential
/// backoff, then reverse-compensate the succeeded prefix via <c>onError</c>
/// targets, then dead-letter unhandled failures to <c>wf:{tenant}:dlq</c>.
///
/// <c>durable</c>/<c>handler</c> modes checkpoint after each transition so the
/// instance can resume. A per-run Redis lease prevents double-dispatch across
/// control-plane replicas. Delivery is at-least-once; duplicate results are
/// ignored because each TaskRun is keyed by (run, task, attempt).
/// </summary>
public class RunLifecycleEngine
{
    private readonly IEventBus _events;
    private readonly ILeaseService _leases;
    private readonly ILogger<RunLifecycleEngine> _logger;

    private readonly string _owner;
    private readonly int _batchSize;
    private readonly TimeSpan _claimIdle;
    private readonly TimeSpan _leaseTtl;

    private readonly int _maxRetries;
    private readonly TimeSpan _retryBaseDelay;
    private readonly TimeSpan _retryMaxDelay;
    private readonly TimeSpan _runTimeout;

    private static readonly JsonSerializerOptions PlanJson = new(JsonSerializerDefaults.Web);

    public RunLifecycleEngine(
        IEventBus events,
        ILeaseService leases,
        IConfiguration config,
        ILogger<RunLifecycleEngine> logger)
    {
        _events = events;
        _leases = leases;
        _logger = logger;

        _owner = $"runs-{Environment.MachineName}-{Guid.NewGuid():N}"[..28];
        _batchSize = ParseInt(config["Workflows:LifecycleBatchSize"], 10);
        _claimIdle = TimeSpan.FromSeconds(ParseInt(config["Workflows:LifecycleClaimIdleSeconds"], 30));
        _leaseTtl = TimeSpan.FromSeconds(ParseInt(config["Workflows:LifecycleLeaseTtlSeconds"], 60));

        _maxRetries = ParseNonNegative(config["Workflows:MaxRetries"], 3);
        _retryBaseDelay = TimeSpan.FromMilliseconds(ParseNonNegative(config["Workflows:RetryBaseDelayMs"], 1000));
        _retryMaxDelay = TimeSpan.FromMilliseconds(ParseNonNegative(config["Workflows:RetryMaxDelayMs"], 30000));
        _runTimeout = TimeSpan.FromSeconds(ParseNonNegative(config["Workflows:RunTimeoutSeconds"], 300));
    }

    // ------------------------------------------------------------------ dispatch

    public async Task<int> DispatchPendingAsync(WorkflowsDbContext db, TaskDispatcher dispatcher, CancellationToken ct)
    {
        var pending = await db.WorkflowRuns
            .Where(r => r.Status == RunStatus.Pending)
            .OrderBy(r => r.Id)
            .Take(_batchSize)
            .ToListAsync(ct);

        var dispatched = 0;
        foreach (var run in pending)
        {
            if (ct.IsCancellationRequested)
                return dispatched;

            if (!await _leases.AcquireAsync(run.TenantId, run.Id.ToString(), _owner, _leaseTtl))
                continue; // another replica owns this run

            try
            {
                if (await DispatchRunAsync(db, dispatcher, run, ct))
                    dispatched++;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Failed to dispatch run {RunId} (tenant {TenantId})", run.Id, run.TenantId);
                await FailRunAsync(db, run, $"dispatch error: {ex.Message}", ct);
            }
        }

        return dispatched;
    }

    private async Task<bool> DispatchRunAsync(WorkflowsDbContext db, TaskDispatcher dispatcher, WorkflowRun run, CancellationToken ct)
    {
        var workflow = await db.Workflows
            .Include(w => w.Tasks)
            .FirstOrDefaultAsync(w => w.Id == run.WorkflowId && w.TenantId == run.TenantId, ct);

        if (workflow == null || workflow.Status != WorkflowStatus.Compiled)
        {
            _logger.LogWarning("Run {RunId} references missing/uncompiled workflow {WorkflowId}; failing it",
                run.Id, run.WorkflowId);
            await FailRunAsync(db, run, "workflow is missing or not compiled", ct);
            return false;
        }

        var startTask = run.StartTaskId != null
            ? workflow.Tasks.FirstOrDefault(t => t.Id == run.StartTaskId)
            : workflow.Tasks.FirstOrDefault(t => t.Key == "root");

        if (startTask == null)
        {
            _logger.LogWarning("Run {RunId} has no start task (StartTaskId {StartTaskId}); failing it",
                run.Id, run.StartTaskId);
            await FailRunAsync(db, run, "start task not found in workflow", ct);
            return false;
        }

        run.Status = RunStatus.Running;
        run.StartedAt = DateTime.UtcNow;
        await db.SaveChangesAsync(ct); // persist running before enqueue so a result can never race a pending run

        await dispatcher.DispatchAsync(run, startTask, workflow.Path, run.InputJson, 1, ct);
        return true;
    }

    // ------------------------------------------------------------------- retries

    /// <summary>
    /// Re-dispatches failed task attempts whose exponential backoff has elapsed.
    /// Called by the hosted service on each tick. The retry is idempotent: it
    /// only fires when the next attempt's TaskRun does not already exist, so a
    /// duplicate tick (or a second replica) cannot double-dispatch the same
    /// attempt.
    ///
    /// A per-TaskRun Redis lease prevents concurrent replicas from
    /// double-dispatching the same retry. This mirrors the per-run lease in
    /// <see cref="DispatchPendingAsync"/>. Each retry is committed individually
    /// so a mid-batch failure never leaves some dispatched but
    /// <c>NextAttemptAt</c> unchanged.
    /// </summary>
    public async Task<int> RetryDueAsync(WorkflowsDbContext db, TaskDispatcher dispatcher, CancellationToken ct)
    {
        var now = DateTime.UtcNow;
        var due = await db.TaskRuns
            .Where(t => t.Status == TaskRunStatus.Failed
                        && t.NextAttemptAt != null
                        && t.NextAttemptAt <= now
                        && !t.IsCompensation)
            .OrderBy(t => t.NextAttemptAt)
            .Take(_batchSize)
            .ToListAsync(ct);

        var retried = 0;
        foreach (var failed in due)
        {
            if (ct.IsCancellationRequested)
                break;

            // Per-TaskRun lease: prevents two replicas from dispatching the
            // same retry concurrently. Uses "taskrun" as the lease namespace
            // (distinct from the "wf:{tenant}:leases:{run}" keys used by
            // DispatchPendingAsync). The key is globally unique per TaskRun.
            if (!await _leases.AcquireAsync("taskrun", failed.Id.ToString(), _owner, _leaseTtl))
                continue;

            try
            {
                var run = await db.WorkflowRuns
                    .Include(r => r.Workflow).ThenInclude(w => w.Tasks)
                    .FirstOrDefaultAsync(r => r.Id == failed.RunId, ct);

                // A stale schedule (the run already finished) is simply cleared.
                if (run == null || run.Status != RunStatus.Running)
                {
                    failed.NextAttemptAt = null;
                    await db.SaveChangesAsync(ct);
                    continue;
                }

                var task = run.Workflow.Tasks.FirstOrDefault(t => t.Id == failed.TaskId);
                if (task == null)
                {
                    failed.NextAttemptAt = null;
                    await db.SaveChangesAsync(ct);
                    continue;
                }

                var nextAttempt = failed.Attempt + 1;
                var already = await db.TaskRuns.AnyAsync(
                    t => t.RunId == failed.RunId && t.TaskId == failed.TaskId && t.Attempt == nextAttempt, ct);
                if (already)
                {
                    failed.NextAttemptAt = null;
                    await db.SaveChangesAsync(ct);
                    continue;
                }

                await dispatcher.DispatchAsync(run, task, run.Workflow.Path, failed.InputJson, nextAttempt, ct);
                failed.NextAttemptAt = null;

                // Commit per-retry: if the next iteration fails, this retry is
                // already dispatched and NextAttemptAt is cleared — no double-dispatch.
                await db.SaveChangesAsync(ct);
                retried++;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Failed to retry TaskRun {TaskRunId} (run {RunId})", failed.Id, failed.RunId);
                // Don't clear NextAttemptAt — the next tick will retry.
            }
            finally
            {
                await _leases.ReleaseAsync("taskrun", failed.Id.ToString(), _owner);
            }
        }

        return retried;
    }

    /// <summary>
    /// Fails runs that have been <c>running</c>/<c>compensating</c> for longer
    /// than <c>RunTimeoutSeconds</c> without completing. This is the safety net
    /// for runs orphaned by a worker outage, a crash mid-chain, or a job queued
    /// before a worker consumer group existed (which Redis Streams never replays).
    /// Without it such runs would show as "running" forever.
    /// </summary>
    public async Task<int> TimeoutStaleRunsAsync(WorkflowsDbContext db, CancellationToken ct)
    {
        if (_runTimeout <= TimeSpan.Zero)
            return 0;

        var cutoff = DateTime.UtcNow - _runTimeout;
        var stale = await db.WorkflowRuns
            .Where(r => (r.Status == RunStatus.Running || r.Status == RunStatus.Compensating)
                        && r.StartedAt != null && r.StartedAt < cutoff)
            .OrderBy(r => r.StartedAt)
            .Take(_batchSize)
            .ToListAsync(ct);

        foreach (var run in stale)
        {
            ct.ThrowIfCancellationRequested();

            // Mark every in-flight task run of the stale run as dead so the run
            // detail is consistent (no dangling "running" tasks under a failed run).
            var inFlight = await db.TaskRuns
                .Where(t => t.RunId == run.Id && t.Status == TaskRunStatus.Running)
                .ToListAsync(ct);
            foreach (var taskRun in inFlight)
            {
                taskRun.Status = TaskRunStatus.Dead;
                taskRun.Error = string.IsNullOrEmpty(taskRun.Error)
                    ? $"timed out after {_runTimeout.TotalSeconds:0}s"
                    : taskRun.Error;
                taskRun.FinishedAt ??= DateTime.UtcNow;
            }

            run.Status = RunStatus.Failed;
            run.Error = $"run timed out after {_runTimeout.TotalSeconds:0}s";
            run.FinishedAt = DateTime.UtcNow;

            _logger.LogWarning(
                "Timed out stale run {RunId} (tenant {TenantId}, started {StartedAt}) after {Seconds}s",
                run.Id, run.TenantId, run.StartedAt, _runTimeout.TotalSeconds);
        }

        if (stale.Count > 0)
            await db.SaveChangesAsync(ct);

        return stale.Count;
    }

    // ------------------------------------------------------------------ results

    public async Task<int> ProcessResultsAsync(
        WorkflowsDbContext db,
        TaskDispatcher dispatcher,
        DurableStateStore store,
        CancellationToken ct)
    {
        var tenants = await db.WorkflowRuns
            .Where(r => r.Status == RunStatus.Running || r.Status == RunStatus.Compensating)
            .Select(r => r.TenantId)
            .Distinct()
            .ToListAsync(ct);

        var processed = 0;
        foreach (var tenant in tenants)
        {
            if (ct.IsCancellationRequested)
                return processed;

            await _events.EnsureGroupAsync(tenant, Streams.Results, ct);

            var messages = await _events.ReadGroupAsync(tenant, Streams.Results, _owner, _batchSize, ct);
            if (messages.Count == 0)
                messages = await _events.ClaimPendingAsync(tenant, Streams.Results, _owner, _claimIdle, _batchSize, ct);

            foreach (var message in messages)
            {
                ct.ThrowIfCancellationRequested();
                try
                {
                    await ProcessResultAsync(db, dispatcher, store, tenant, message, ct);
                }
                catch (Exception ex)
                {
                    _logger.LogError(ex, "Failed to process result {MessageId} for tenant {TenantId}", message.Id, tenant);
                }
                finally
                {
                    await _events.AckAsync(tenant, Streams.Results, message.Id, ct);
                }
                processed++;
            }
        }

        return processed;
    }

    private async Task ProcessResultAsync(
        WorkflowsDbContext db,
        TaskDispatcher dispatcher,
        DurableStateStore store,
        string tenant,
        StreamMessage message,
        CancellationToken ct)
    {
        TaskResult result;
        try
        {
            result = TaskResult.FromFields(message.Fields);
        }
        catch (Exception ex)
        {
            _logger.LogWarning("Malformed result {MessageId}: {Error}", message.Id, ex.Message);
            return;
        }

        if (!string.Equals(result.TenantId, tenant, StringComparison.Ordinal)
            || !Guid.TryParse(result.RunId, out var runId)
            || !Guid.TryParse(result.TaskId, out var taskId))
        {
            return;
        }

        var taskRun = await db.TaskRuns
            .FirstOrDefaultAsync(t => t.RunId == runId && t.TaskId == taskId && t.Attempt == result.Attempt, ct);

        if (taskRun == null)
        {
            _logger.LogDebug("Ignoring orphan result for run {RunId} task {TaskId} attempt {Attempt}",
                runId, taskId, result.Attempt);
            return;
        }

        if (taskRun.Status is TaskRunStatus.Succeeded or TaskRunStatus.Failed
            or TaskRunStatus.Compensated or TaskRunStatus.Dead)
        {
            _logger.LogDebug("Ignoring duplicate result for taskRun {TaskRunId}", taskRun.Id);
            return;
        }

        var run = await db.WorkflowRuns
            .Include(r => r.Workflow).ThenInclude(w => w.Tasks)
            .FirstOrDefaultAsync(r => r.Id == runId, ct);

        if (run == null || run.Status is not (RunStatus.Running or RunStatus.Compensating))
        {
            _logger.LogDebug("Ignoring result for non-running run {RunId} (status {Status})", runId, run?.Status);
            return;
        }

        var task = run.Workflow.Tasks.FirstOrDefault(t => t.Id == taskId);
        if (task == null)
            return;

        var now = DateTime.UtcNow;

        // A result that arrives while the run is compensating is a compensation
        // step's completion, not a normal chain advance.
        if (run.Status == RunStatus.Compensating)
        {
            await HandleCompensationResultAsync(db, dispatcher, store, run, task, taskRun, result, ct);
            await db.SaveChangesAsync(ct);
            return;
        }

        if (result.Success)
        {
            taskRun.Status = TaskRunStatus.Succeeded;
            taskRun.OutputJson = result.Output;
            taskRun.FinishedAt = now;
            await HandleTaskSucceededAsync(db, dispatcher, store, run, task, taskRun, ct);
        }
        else
        {
            taskRun.Status = TaskRunStatus.Failed;
            taskRun.Error = result.Error;
            taskRun.FinishedAt = now;
            await HandleTaskFailedAsync(db, dispatcher, store, run, task, taskRun, ct);
        }

        await db.SaveChangesAsync(ct);
    }

    private async Task HandleTaskSucceededAsync(
        WorkflowsDbContext db,
        TaskDispatcher dispatcher,
        DurableStateStore store,
        WorkflowRun run,
        WorkflowTask task,
        TaskRun taskRun,
        CancellationToken ct)
    {
        if (task.NextId == null)
        {
            // Terminal task on the success chain — the run is complete.
            run.Status = RunStatus.Succeeded;
            run.OutputJson = taskRun.OutputJson;
            run.FinishedAt = DateTime.UtcNow;
            if (run.Workflow.Mode != WorkflowMode.Function)
                await store.ClearAsync(InstanceIdFor(run), ct);
            await _leases.ReleaseAsync(run.TenantId, run.Id.ToString(), _owner);
            return;
        }

        var nextTask = run.Workflow.Tasks.FirstOrDefault(t => t.Id == task.NextId);
        if (nextTask == null)
        {
            await FailRunAsync(db, run, $"success edge target {task.NextId} is missing from the compiled workflow", ct);
            return;
        }

        // Durable/handler modes checkpoint before each transition so the instance
        // can resume from this exact point (next task + its input).
        if (run.Workflow.Mode != WorkflowMode.Function)
            await store.CheckpointAsync(InstanceIdFor(run), nextTask.Id, taskRun.OutputJson, run.CorrelationId, ct);

        await dispatcher.DispatchAsync(run, nextTask, run.Workflow.Path, taskRun.OutputJson, 1, ct);
    }

    private async Task HandleTaskFailedAsync(
        WorkflowsDbContext db,
        TaskDispatcher dispatcher,
        DurableStateStore store,
        WorkflowRun run,
        WorkflowTask task,
        TaskRun taskRun,
        CancellationToken ct)
    {
        // 1. Retry with exponential backoff before compensating.
        if (taskRun.Attempt <= _maxRetries)
        {
            taskRun.NextAttemptAt = DateTime.UtcNow.Add(BackoffFor(taskRun.Attempt));
            _logger.LogInformation(
                "Task {TaskKey} attempt {Attempt} failed; retrying at {When} (run {RunId})",
                task.Key, taskRun.Attempt, taskRun.NextAttemptAt, run.Id);
            return; // run stays Running; RetryDueAsync re-dispatches once due.
        }

        // 2. Retries exhausted — compensate, else dead-letter.
        var plan = await BuildCompensationPlanAsync(db, run, task, taskRun, ct);
        if (plan.Count > 0)
        {
            run.Status = RunStatus.Compensating;
            run.Error = taskRun.Error;
            run.CompensationPlanJson = SerializePlan(plan);
            await DispatchCompensationAsync(db, dispatcher, run, plan[0], ct);
            return;
        }

        // 3. Unhandled: dead-letter and mark the run dead.
        taskRun.Status = TaskRunStatus.Dead;
        await dispatcher.DeadLetterTaskAsync(run, task, taskRun, ct);
        run.Status = RunStatus.Dead;
        run.Error = taskRun.Error;
        run.FinishedAt = DateTime.UtcNow;
        if (run.Workflow.Mode != WorkflowMode.Function)
            await store.ClearAsync(InstanceIdFor(run), ct);
        await _leases.ReleaseAsync(run.TenantId, run.Id.ToString(), _owner);
    }

    /// <summary>
    /// Builds the ordered saga compensation plan:
    ///   <list type="number">
    ///   <item>The failing task's own <c>onError</c> edge (its error handler).</item>
    ///   <item>Each already-succeeded task's <c>onError</c> target, walked in
    ///   reverse completion order to undo side effects.</item>
    ///   </list>
    /// Duplicate targets are run once (first occurrence wins).
    /// </summary>
    private async Task<List<CompensationStep>> BuildCompensationPlanAsync(
        WorkflowsDbContext db,
        WorkflowRun run,
        WorkflowTask failingTask,
        TaskRun failedRun,
        CancellationToken ct)
    {
        var plan = new List<CompensationStep>();
        var seen = new HashSet<Guid>();

        if (failingTask.OnErrorId is Guid ownError && seen.Add(ownError))
            plan.Add(new CompensationStep(ownError, failedRun.InputJson));

        var succeeded = await db.TaskRuns
            .Where(t => t.RunId == run.Id && t.Status == TaskRunStatus.Succeeded && !t.IsCompensation)
            .OrderByDescending(t => t.FinishedAt)
            .ToListAsync(ct);

        foreach (var tr in succeeded)
        {
            var task = run.Workflow.Tasks.FirstOrDefault(t => t.Id == tr.TaskId);
            if (task?.OnErrorId is Guid target && seen.Add(target))
                plan.Add(new CompensationStep(target, tr.OutputJson));
        }

        return plan;
    }

    private async Task HandleCompensationResultAsync(
        WorkflowsDbContext db,
        TaskDispatcher dispatcher,
        DurableStateStore store,
        WorkflowRun run,
        WorkflowTask task,
        TaskRun taskRun,
        TaskResult result,
        CancellationToken ct)
    {
        var plan = DeserializePlan(run.CompensationPlanJson);

        if (result.Success)
        {
            taskRun.Status = TaskRunStatus.Compensated;
            taskRun.OutputJson = result.Output;
            taskRun.FinishedAt = DateTime.UtcNow;

            if (plan.Count > 0)
                plan.RemoveAt(0); // the in-flight step just completed

            if (plan.Count > 0)
            {
                run.CompensationPlanJson = SerializePlan(plan);
                await DispatchCompensationAsync(db, dispatcher, run, plan[0], ct);
                return;
            }

            // All compensations done — the run is compensated but still failed.
            run.CompensationPlanJson = null;
            run.Status = RunStatus.Failed;
            run.FinishedAt = DateTime.UtcNow;
            if (run.Workflow.Mode != WorkflowMode.Function)
                await store.ClearAsync(InstanceIdFor(run), ct);
            await _leases.ReleaseAsync(run.TenantId, run.Id.ToString(), _owner);
            return;
        }

        // A compensation step that fails is itself dead-lettered; the run is dead.
        taskRun.Status = TaskRunStatus.Dead;
        taskRun.Error = result.Error;
        taskRun.FinishedAt = DateTime.UtcNow;
        await dispatcher.DeadLetterTaskAsync(run, task, taskRun, ct);

        run.CompensationPlanJson = null;
        run.Status = RunStatus.Dead;
        run.Error = taskRun.Error;
        run.FinishedAt = DateTime.UtcNow;
        if (run.Workflow.Mode != WorkflowMode.Function)
            await store.ClearAsync(InstanceIdFor(run), ct);
        await _leases.ReleaseAsync(run.TenantId, run.Id.ToString(), _owner);
    }

    private async Task DispatchCompensationAsync(
        WorkflowsDbContext db,
        TaskDispatcher dispatcher,
        WorkflowRun run,
        CompensationStep step,
        CancellationToken ct)
    {
        var task = run.Workflow.Tasks.FirstOrDefault(t => t.Id == step.TaskId);
        if (task == null)
        {
            // Dangling compensation reference (should be impossible post-compile,
            // but a hand-edited DB could break it): skip straight to the next step.
            _logger.LogWarning("Compensation target {TaskId} missing from workflow {WorkflowId}", step.TaskId, run.WorkflowId);
            var plan = DeserializePlan(run.CompensationPlanJson);
            if (plan.Count > 0)
                plan.RemoveAt(0);
            run.CompensationPlanJson = plan.Count > 0 ? SerializePlan(plan) : null;
            return;
        }

        await dispatcher.DispatchAsync(run, task, run.Workflow.Path, step.Input, 1, ct, isCompensation: true);
    }

    private TimeSpan BackoffFor(int failedAttempt)
    {
        if (_retryBaseDelay <= TimeSpan.Zero)
            return TimeSpan.Zero;
        var ms = _retryBaseDelay.TotalMilliseconds * Math.Pow(2, failedAttempt - 1);
        return TimeSpan.FromMilliseconds(Math.Min(ms, _retryMaxDelay.TotalMilliseconds));
    }

    private static List<CompensationStep> DeserializePlan(string? planJson)
        => planJson == null
            ? new List<CompensationStep>()
            : JsonSerializer.Deserialize<List<CompensationStep>>(planJson, PlanJson) ?? new List<CompensationStep>();

    private static string SerializePlan(List<CompensationStep> plan)
        => JsonSerializer.Serialize(plan, PlanJson);

    private sealed record CompensationStep(Guid TaskId, string? Input);

    private async Task FailRunAsync(WorkflowsDbContext db, WorkflowRun run, string error, CancellationToken ct)
    {
        run.Status = RunStatus.Failed;
        run.Error = error;
        run.FinishedAt = DateTime.UtcNow;
        await db.SaveChangesAsync(ct);
        await _leases.ReleaseAsync(run.TenantId, run.Id.ToString(), _owner);
    }

    private static string InstanceIdFor(WorkflowRun run)
        => run.Workflow.Mode == WorkflowMode.Handler
            ? $"handler:{run.WorkflowId}"
            : run.Id.ToString();

    private static int ParseInt(string? text, int fallback)
        => int.TryParse(text, out var value) && value > 0 ? value : fallback;

    private static int ParseNonNegative(string? text, int fallback)
        => int.TryParse(text, out var value) && value >= 0 ? value : fallback;
}