"""Running workflows, resumably.

The engine executes a definition step by step, persisting after each one. A run
interrupted anywhere can be resumed, and resuming never repeats a step that
already succeeded — which is what makes retries safe on a workflow that files
documents and raises invoices.

The engine owns sequencing. Tools never call each other, precisely so that the
full graph is visible here and can be recorded, paused and resumed.
"""

from __future__ import annotations

from collections.abc import Callable
from copy import deepcopy
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any, Protocol, runtime_checkable

from app.tools.base import ExecutionPolicy, ToolResult
from app.tools.registry import ToolRegistry
from app.workflow.state import RunState, StepRecord, StepState, WorkflowRun


class ConcurrentUpdateError(RuntimeError):
    """Something else changed this run since it was read.

    Raised rather than returned, because a lost update is a bug in whatever is
    running two writers over one run — not an outcome a caller should be able to
    ignore by forgetting to check a boolean.
    """


@runtime_checkable
class RunStore(Protocol):
    """Where runs are persisted between steps."""

    def save(self, run: WorkflowRun) -> None: ...

    def load(self, run_id: str) -> WorkflowRun | None: ...

    def resumable(self) -> list[WorkflowRun]: ...

    def archive(self, older_than: datetime) -> int: ...


class InMemoryRunStore:
    """Runs held in process — for tests, and for a deployment before Postgres.

    Honest about its limit: nothing survives a restart, which is exactly the
    situation resumability exists for. Production uses PostgresRunStore; the port
    is here so swapping it is a configuration change.

    It enforces the same version check as the durable store. That is not
    ceremony: a store with no concurrency semantics makes every test pass that
    would fail against Postgres, and the difference would surface first in
    production.
    """

    def __init__(self) -> None:
        self._runs: dict[str, WorkflowRun] = {}
        self._versions: dict[str, int] = {}

    def save(self, run: WorkflowRun) -> None:
        stored = self._versions.get(run.id, 0)

        if stored != run.version:
            raise ConcurrentUpdateError(
                f"Run {run.id} was version {run.version} here and {stored} in the store."
            )

        run.touch()
        run.version += 1
        self._versions[run.id] = run.version
        # Copied so a caller mutating its object after saving does not silently
        # change what the store believes it holds — which a database never does.
        self._runs[run.id] = deepcopy(run)

    def load(self, run_id: str) -> WorkflowRun | None:
        stored = self._runs.get(run_id)

        return deepcopy(stored) if stored is not None else None

    def resumable(self) -> list[WorkflowRun]:
        return [
            deepcopy(r)
            for r in self._runs.values()
            if r.state.is_resumable and r.archived_at is None
        ]

    def archive(self, older_than: datetime) -> int:
        archived = 0

        for run in self._runs.values():
            if run.state.is_terminal and run.archived_at is None and run.updated_at < older_than:
                run.archived_at = datetime.now(UTC)
                archived += 1

        return archived

    def __len__(self) -> int:
        return len(self._runs)


@dataclass(frozen=True, slots=True)
class Step:
    """One step of a definition.

    ``build_input`` turns the accumulated run context into the Tool's arguments,
    so a step depends on earlier results without any Tool knowing it is part of
    a workflow.
    """

    name: str
    tool: str
    build_input: Callable[[dict[str, Any]], dict[str, Any]] = lambda ctx: {}
    optional: bool = False
    """An optional step that fails is skipped and the run continues. For work
    that improves an outcome without being required for it."""

    skip_when: Callable[[dict[str, Any]], bool] | None = None
    """Decided before the step runs, from what earlier steps produced.

    Distinct from ``optional``, which is about a step that ran and failed. This
    is about a step that should not run at all: the filing step skips itself
    when the document was read and named nobody, so the run ends without a
    proposal and the document can be asked about instead of reaching a reviewer
    with an empty client field.
    """

    max_attempts: int = 3


@dataclass(frozen=True, slots=True)
class Workflow:
    """A named sequence of steps."""

    name: str
    steps: tuple[Step, ...]

    def new_run(self, context: dict[str, Any] | None = None) -> WorkflowRun:
        run = WorkflowRun(
            workflow=self.name,
            context=dict(context or {}),
            steps=[StepRecord(name=s.name, tool=s.tool) for s in self.steps],
        )

        # Published so a step can identify its own run — an idempotency key and
        # a memory note both need it, and reaching for it any other way would
        # mean handing every Tool a reference to the run it belongs to.
        run.context.setdefault("run_id", run.id)

        return run


class WorkflowEngine:
    """Executes workflows against a registry of tools."""

    def __init__(self, tools: ToolRegistry, store: RunStore | None = None) -> None:
        self._tools = tools
        # `store or InMemoryRunStore()` is wrong here and was a real bug: a store
        # that implements __len__ is FALSY while empty, so a caller's brand-new
        # store was silently discarded and every save went into an orphan the
        # caller could not see. Identity check, not truthiness.
        self._store = store if store is not None else InMemoryRunStore()

    @property
    def store(self) -> RunStore:
        return self._store

    def start(self, workflow: Workflow, context: dict[str, Any] | None = None) -> WorkflowRun:
        run = workflow.new_run(context)
        self._store.save(run)

        return self.resume(workflow, run)

    def resume(self, workflow: Workflow, run: WorkflowRun) -> WorkflowRun:
        """Continue a run from wherever it stopped.

        Steps already done are not re-executed — not skipped conditionally, but
        never reached, because the loop takes the first step that is not done.
        A resumed run that re-ran a completed 'raise invoice' step would bill a
        client twice.
        """
        if run.state.is_terminal:
            return run

        by_name = {s.name: s for s in workflow.steps}
        run.state = RunState.RUNNING

        while (record := run.next_step) is not None:
            step = by_name.get(record.name)

            if step is None:
                # The definition changed under a stored run. Failing loudly beats
                # guessing which step was meant.
                return self._fail(run, f"Step '{record.name}' is no longer part of {workflow.name}.")

            outcome = self._run_step(step, record, run)
            self._store.save(run)

            if outcome is RunState.AWAITING_APPROVAL:
                run.state = RunState.AWAITING_APPROVAL
                self._store.save(run)

                return run

            if outcome is RunState.FAILED:
                return self._fail(run, record.error or f"Step '{record.name}' failed.")

        run.state = RunState.COMPLETED
        self._store.save(run)

        return run

    def approve(
        self,
        workflow: Workflow,
        run: WorkflowRun,
        *,
        approved_by: str | None = None,
        amendments: dict[str, Any] | None = None,
    ) -> WorkflowRun:
        """Perform the paused step, now that a human has authorised it.

        Approval **executes** the step. An earlier version marked it succeeded
        and resumed, which meant an approved upload never actually uploaded:
        the run reported ``completed`` while nothing had been written. Anything
        that stops short of running the tool is a lie told to the person who
        clicked approve.

        ``amendments`` override the proposed arguments. This is what makes the
        approval queue useful rather than a yes/no box — the commonest correction
        is the client a document was about to be filed against, and a human who
        can only reject has to start the whole run again to fix a single id.
        """
        if run.state is not RunState.AWAITING_APPROVAL:
            raise ValueError(f"Run {run.id} is {run.state}, not awaiting approval.")

        record = next((r for r in run.steps if r.state is StepState.AWAITING_APPROVAL), None)

        if record is None:
            return self._fail(run, "Run is awaiting approval but no step is paused.")

        record.approved = True
        record.approved_by = approved_by
        record.amendments = dict(amendments or {})
        record.error = None

        # Returned to PENDING so the ordinary loop picks it up and runs it. The
        # approval is recorded on the step, not spent on this one call, so a
        # retry after a transient CMS failure does not ask the human again for
        # something they have already authorised.
        record.state = StepState.PENDING

        self._store.save(run)

        return self.resume(workflow, run)

    def apply_decision(
        self,
        workflow: Workflow,
        run: WorkflowRun,
        *,
        accepted: bool,
        data: dict[str, Any] | None = None,
        reason: str | None = None,
        redo_from: str | None = None,
    ) -> WorkflowRun:
        """Record what somebody else decided about a paused step.

        The counterpart to ``approve`` for work the CMS performs. The step is not
        re-executed — the filing already happened over there — so its result is
        written in directly and the run continues.

        ``redo_from`` handles a reviewer sending the work back: every step from
        that one onward returns to PENDING and runs again, which is what "read
        the document again" actually means. Steps before it keep their results,
        because nothing about them was questioned.
        """
        record = next((r for r in run.steps if r.state is StepState.AWAITING_APPROVAL), None)

        if record is None:
            return self._fail(run, "Run is awaiting a decision but no step is paused.")

        if accepted:
            record.state = StepState.SUCCEEDED
            record.error = None
            record.finished_at = datetime.now(UTC)
            record.output = {**record.output, **(data or {})}
            run.context[record.name] = {**record.output, "approved": True}

            self._store.save(run)

            return self.resume(workflow, run)

        if redo_from is not None:
            return self._redo(workflow, run, redo_from, reason)

        # Refused outright. A rejection is a decision, not a malfunction — but it
        # is terminal, and recording it as anything else would leave the run
        # looking like it might still complete.
        record.state = StepState.FAILED
        record.error = reason or "Rejected by a reviewer."
        record.finished_at = datetime.now(UTC)

        return self._fail(run, record.error)

    # ── Internals ─────────────────────────────────────────────────────────

    def _redo(
        self, workflow: Workflow, run: WorkflowRun, from_step: str, reason: str | None
    ) -> WorkflowRun:
        """Rewind to a step and run everything from there again."""
        names = [s.name for s in workflow.steps]

        if from_step not in names:
            return self._fail(run, f"Cannot redo from '{from_step}': no such step.")

        for record in run.steps[names.index(from_step):]:
            record.state = StepState.PENDING
            record.error = None
            record.output = {}
            record.attempts = 0
            # Cleared so a rebuilt input cannot silently read a stale value that
            # the reviewer's whole point was to have recomputed.
            run.context.pop(record.name, None)

        # Counted so a resubmission can carry a new idempotency key — the CMS
        # returns the same proposal for a repeated one, which would hand the
        # reviewer back the answer they just rejected.
        run.context["attempt"] = int(run.context.get("attempt", 1)) + 1
        run.context["redo_reason"] = reason

        self._store.save(run)

        return self.resume(workflow, run)

    def _run_step(self, step: Step, record: StepRecord, run: WorkflowRun) -> RunState:
        # Before the attempt is counted: a step that was never meant to run has
        # not attempted anything, and recording an attempt would make the
        # retry figures read as work that failed.
        if step.skip_when is not None:
            try:
                skip = bool(step.skip_when(run.context))
            except Exception:  # noqa: BLE001
                # A predicate that cannot decide is not grounds to skip. Running
                # the step is the behaviour that existed before this feature.
                skip = False

            if skip:
                record.started_at = datetime.now(UTC)

                return self._skip(record)

        record.attempts += 1
        record.started_at = datetime.now(UTC)

        if not self._tools.has(step.tool):
            record.state = StepState.FAILED
            record.error = f"No tool named '{step.tool}'."

            return RunState.FAILED if not step.optional else self._skip(record)

        tool = self._tools.get(step.tool)

        try:
            arguments = step.build_input(run.context)
        except Exception as exc:  # noqa: BLE001 - a bad input builder is a step failure
            record.state = StepState.FAILED
            record.error = f"Could not build input: {exc}"

            return RunState.FAILED if not step.optional else self._skip(record)

        if tool.policy is ExecutionPolicy.REQUIRES_APPROVAL and record.approved:
            # The one path that performs an approved write. Amendments last:
            # a human correcting the client id must beat what the AI proposed,
            # and build_input has just recomputed the rest from fresh context.
            result: ToolResult = tool.execute_approved(**{**arguments, **record.amendments})
        else:
            result = tool.execute(**arguments)

        record.duration_ms = result.duration_ms
        record.confidence = result.confidence
        record.finished_at = datetime.now(UTC)

        # An approval-gated tool returns a proposal rather than acting. That is
        # a pause, not a failure, and must not consume retry attempts.
        if tool.policy is ExecutionPolicy.REQUIRES_APPROVAL and result.error == "approval_required":
            record.state = StepState.AWAITING_APPROVAL
            record.output = result.data
            record.proposed = dict(result.data.get("proposed") or {})
            record.attempts -= 1

            return RunState.AWAITING_APPROVAL

        # A step whose outcome is decided elsewhere — a proposal submitted to the
        # CMS, waiting on a human. Also a pause, for the same reason: retrying
        # would submit it again, and continuing would pretend it was done.
        if result.error == ToolResult.AWAITING_DECISION:
            record.state = StepState.AWAITING_APPROVAL
            record.output = result.data
            record.attempts -= 1

            return RunState.AWAITING_APPROVAL

        if result.ok:
            record.state = StepState.SUCCEEDED
            record.output = result.data
            record.error = None

            # When each step actually finished, so a later step can report the
            # real sequence rather than stamping everything with the moment it
            # got round to sending them.
            run.context.setdefault("step_times", {})[step.name] = record.finished_at.isoformat()

            # And how long it took. Recorded on the step already, but only in this
            # agent's own database — which the CMS cannot query, so "OCR takes
            # ninety seconds and classification takes forty milliseconds" was a
            # fact nobody on the CMS side could establish (ADR-0011).
            run.context.setdefault("step_durations", {})[step.name] = record.duration_ms
            # Published under the step's name so a later step can reach it. An
            # approved step is flagged so a later step (and the audit trail) can
            # tell autonomous work from authorised work.
            #
            # The confidence travels with it. A tool reports that alongside its
            # data rather than inside it, so it landed on the step record and
            # went no further — and the proposal a reviewer sees reads its
            # confidence out of the context. Every document therefore arrived
            # with none, scored high risk on the "we do not know how sure it
            # was" branch, and showed "—" in the queue's confidence column. The
            # classifier had been perfectly sure; nothing downstream could tell.
            #
            # setdefault rather than an overwrite: a tool that puts a confidence
            # inside its own data means something more specific by it.
            published = dict(result.data)
            published.setdefault("confidence", result.confidence)

            if record.approved:
                published["approved"] = True

            run.context[step.name] = published

            return RunState.RUNNING

        record.error = result.error

        if record.attempts < step.max_attempts:
            # Left PENDING so the loop picks it up again. The caller decides
            # when to retry; the engine does not sleep or spin.
            record.state = StepState.PENDING

            return RunState.RUNNING

        record.state = StepState.FAILED

        return self._skip(record) if step.optional else RunState.FAILED

    @staticmethod
    def _skip(record: StepRecord) -> RunState:
        record.state = StepState.SKIPPED

        return RunState.RUNNING

    def _fail(self, run: WorkflowRun, error: str) -> WorkflowRun:
        run.state = RunState.FAILED
        run.error = error
        self._store.save(run)

        return run