"""The workflow engine.

Nearly every test here is about interruption. A workflow that only works when
nothing goes wrong is not a workflow engine — the process will restart mid-run,
a step will time out, the CMS will be briefly unreachable, and none of that may
cost the work already done or cause it to be repeated.
"""

from __future__ import annotations

import pytest

from app.tools.base import ExecutionPolicy, Tool, ToolResult
from app.tools.registry import ToolRegistry
from app.workflow.engine import InMemoryRunStore, Step, Workflow, WorkflowEngine
from app.workflow.state import RunState, StepState


def make_tool(name: str, *, policy=ExecutionPolicy.AUTOMATIC, behaviour=None) -> Tool:
    """Build a Tool whose behaviour a test controls."""
    calls: list[dict] = []

    def run(self, **kwargs):
        calls.append(kwargs)
        if behaviour is not None:
            return behaviour(len(calls), kwargs)
        return ToolResult.success({"ran": name, "call": len(calls)})

    cls = type(
        f"Tool_{name}",
        (Tool,),
        {"name": name, "description": f"{name} tool", "policy": policy, "run": run},
    )

    tool = cls()
    tool.calls = calls  # type: ignore[attr-defined]

    return tool


@pytest.fixture
def registry() -> ToolRegistry:
    return ToolRegistry()


class TestHappyPath:
    def test_a_run_completes_every_step(self, registry):
        registry.register(make_tool("first"))
        registry.register(make_tool("second"))

        workflow = Workflow("simple", (Step("a", "first"), Step("b", "second")))
        run = WorkflowEngine(registry).start(workflow)

        assert run.state is RunState.COMPLETED
        assert run.progress == "2/2"

    def test_a_step_can_use_an_earlier_step_output(self, registry):
        registry.register(make_tool("producer", behaviour=lambda n, kw: ToolResult.success({"value": 42})))
        registry.register(make_tool("consumer"))

        workflow = Workflow(
            "chained",
            (
                Step("produce", "producer"),
                Step("consume", "consumer", build_input=lambda ctx: {"got": ctx["produce"]["value"]}),
            ),
        )

        WorkflowEngine(registry).start(workflow)

        assert registry.get("consumer").calls[0] == {"got": 42}


class TestResumability:
    def test_resuming_never_re_runs_a_completed_step(self, registry):
        """The property the whole engine exists for.

        A resumed run that re-ran a completed 'raise invoice' step would bill a
        client twice.
        """
        registry.register(make_tool("charge"))
        registry.register(
            make_tool("flaky", behaviour=lambda n, kw: ToolResult.failure("network") if n == 1 else ToolResult.success())
        )

        workflow = Workflow(
            "billing",
            (Step("charge", "charge"), Step("notify", "flaky", max_attempts=1)),
        )

        engine = WorkflowEngine(registry)
        run = engine.start(workflow)

        assert run.state is RunState.FAILED
        assert len(registry.get("charge").calls) == 1

        # Reset the failed step the way an operator retrying would, then resume.
        run.step("notify").state = StepState.PENDING
        run.step("notify").attempts = 0
        run.state = RunState.RUNNING
        run = engine.resume(workflow, run)

        assert run.state is RunState.COMPLETED
        assert len(registry.get("charge").calls) == 1, "The completed step ran a second time."

    def test_a_run_survives_being_reloaded_from_the_store(self, registry):
        registry.register(make_tool("first"))
        registry.register(make_tool("second", behaviour=lambda n, kw: ToolResult.failure("down")))

        workflow = Workflow(
            "interrupted", (Step("a", "first"), Step("b", "second", max_attempts=1))
        )
        store = InMemoryRunStore()
        run = WorkflowEngine(registry, store).start(workflow)

        # A fresh engine, as after a restart.
        reloaded = store.load(run.id)

        assert reloaded is not None
        assert reloaded.step("a").state is StepState.SUCCEEDED
        assert reloaded.context["a"]["ran"] == "first"

    def test_the_store_lists_runs_that_can_be_resumed(self, registry):
        registry.register(make_tool("pauser", policy=ExecutionPolicy.REQUIRES_APPROVAL))

        workflow = Workflow("paused", (Step("p", "pauser"),))
        store = InMemoryRunStore()
        WorkflowEngine(registry, store).start(workflow)

        assert len(store.resumable()) == 1

    def test_a_completed_run_is_not_resumed_again(self, registry):
        registry.register(make_tool("once"))

        workflow = Workflow("done", (Step("a", "once"),))
        engine = WorkflowEngine(registry)
        run = engine.start(workflow)

        engine.resume(workflow, run)

        assert len(registry.get("once").calls) == 1


class TestRetries:
    def test_a_step_is_retried_up_to_its_limit(self, registry):
        registry.register(
            make_tool("unstable", behaviour=lambda n, kw: ToolResult.failure("boom") if n < 3 else ToolResult.success())
        )

        workflow = Workflow("retrying", (Step("a", "unstable", max_attempts=3),))
        run = WorkflowEngine(registry).start(workflow)

        assert run.state is RunState.COMPLETED
        assert run.step("a").attempts == 3

    def test_exhausting_attempts_fails_the_run(self, registry):
        registry.register(make_tool("always_fails", behaviour=lambda n, kw: ToolResult.failure("nope")))

        workflow = Workflow("doomed", (Step("a", "always_fails", max_attempts=2),))
        run = WorkflowEngine(registry).start(workflow)

        assert run.state is RunState.FAILED
        assert run.step("a").attempts == 2
        assert "nope" in run.error

    def test_an_optional_step_is_skipped_rather_than_failing_the_run(self, registry):
        registry.register(make_tool("nice_to_have", behaviour=lambda n, kw: ToolResult.failure("unavailable")))
        registry.register(make_tool("essential"))

        workflow = Workflow(
            "tolerant",
            (
                Step("extra", "nice_to_have", optional=True, max_attempts=1),
                Step("core", "essential"),
            ),
        )
        run = WorkflowEngine(registry).start(workflow)

        assert run.state is RunState.COMPLETED
        assert run.step("extra").state is StepState.SKIPPED

    def test_a_missing_tool_fails_the_run_clearly(self, registry):
        workflow = Workflow("broken", (Step("a", "does_not_exist"),))
        run = WorkflowEngine(registry).start(workflow)

        assert run.state is RunState.FAILED
        assert "does_not_exist" in run.step("a").error

    def test_a_failing_input_builder_is_a_step_failure_not_a_crash(self, registry):
        registry.register(make_tool("fine"))

        workflow = Workflow(
            "bad_input",
            (Step("a", "fine", build_input=lambda ctx: {"x": ctx["missing"]}, max_attempts=1),),
        )
        run = WorkflowEngine(registry).start(workflow)

        assert run.state is RunState.FAILED
        assert "Could not build input" in run.step("a").error


class TestApproval:
    def test_an_approval_step_pauses_rather_than_failing(self, registry):
        registry.register(make_tool("writes", policy=ExecutionPolicy.REQUIRES_APPROVAL))

        workflow = Workflow("gated", (Step("w", "writes"),))
        run = WorkflowEngine(registry).start(workflow)

        # ADR-0004: the AI proposes, a human approves. A pause is the expected
        # outcome, not an error.
        assert run.state is RunState.AWAITING_APPROVAL
        assert run.step("w").state is StepState.AWAITING_APPROVAL

    def test_pausing_does_not_consume_retry_attempts(self, registry):
        registry.register(make_tool("writes", policy=ExecutionPolicy.REQUIRES_APPROVAL))

        workflow = Workflow("gated", (Step("w", "writes", max_attempts=2),))
        run = WorkflowEngine(registry).start(workflow)

        # Otherwise a workflow awaiting a human overnight would exhaust its
        # retries and fail on approval. Waiting is not trying: the step has not
        # been attempted, so the full budget is still there for the real
        # execution that approval triggers.
        assert run.step("w").attempts == 0
        assert run.state is RunState.AWAITING_APPROVAL

    def test_the_proposal_is_carried_for_the_human_to_review(self, registry):
        registry.register(make_tool("writes", policy=ExecutionPolicy.REQUIRES_APPROVAL))

        workflow = Workflow(
            "gated", (Step("w", "writes", build_input=lambda ctx: {"client_id": 7, "city": "Lahore"}),)
        )
        run = WorkflowEngine(registry).start(workflow)

        assert run.step("w").output["proposed"] == {"client_id": 7, "city": "Lahore"}

    def test_approving_continues_the_run(self, registry):
        registry.register(make_tool("writes", policy=ExecutionPolicy.REQUIRES_APPROVAL))
        registry.register(make_tool("after"))

        workflow = Workflow("gated", (Step("w", "writes"), Step("a", "after")))
        engine = WorkflowEngine(registry)
        run = engine.start(workflow)

        run = engine.approve(workflow, run)

        assert run.state is RunState.COMPLETED
        assert run.context["w"]["approved"] is True

    def test_approving_a_run_that_is_not_waiting_is_refused(self, registry):
        registry.register(make_tool("plain"))

        workflow = Workflow("simple", (Step("a", "plain"),))
        engine = WorkflowEngine(registry)
        run = engine.start(workflow)

        with pytest.raises(ValueError, match="not awaiting approval"):
            engine.approve(workflow, run)


class TestApprovalActuallyWrites:
    """Approval must *perform* the work, not merely tick it off.

    An earlier engine marked the paused step succeeded and resumed, so an
    approved upload never uploaded: the run reported COMPLETED while nothing
    had been written. Every test above passed throughout — they all checked run
    state and never asked whether the tool ran. These ask.
    """

    def test_the_tool_does_not_run_while_the_step_is_only_proposed(self, registry):
        tool = make_tool("writes", policy=ExecutionPolicy.REQUIRES_APPROVAL)
        registry.register(tool)

        WorkflowEngine(registry).start(Workflow("gated", (Step("w", "writes"),)))

        # The whole point of the gate: proposing is not doing.
        assert tool.calls == []

    def test_approving_executes_the_tool(self, registry):
        tool = make_tool("writes", policy=ExecutionPolicy.REQUIRES_APPROVAL)
        registry.register(tool)

        workflow = Workflow(
            "gated", (Step("w", "writes", build_input=lambda ctx: {"client_id": 7}),)
        )
        engine = WorkflowEngine(registry)
        run = engine.approve(workflow, engine.start(workflow))

        assert run.state is RunState.COMPLETED
        assert tool.calls == [{"client_id": 7}]

    def test_a_human_can_correct_the_proposal_before_approving(self, registry):
        tool = make_tool("writes", policy=ExecutionPolicy.REQUIRES_APPROVAL)
        registry.register(tool)

        workflow = Workflow(
            "gated",
            (Step("w", "writes", build_input=lambda ctx: {"client_id": 7, "type": "cnic"}),),
        )
        engine = WorkflowEngine(registry)
        run = engine.start(workflow)

        # The commonest correction by far: the AI matched the wrong client. A
        # reviewer who can only accept or reject has to discard the whole run to
        # fix one id.
        run = engine.approve(workflow, run, amendments={"client_id": 99})

        assert tool.calls == [{"client_id": 99, "type": "cnic"}]

    def test_the_correction_is_visible_in_the_audit_trail(self, registry):
        registry.register(make_tool("writes", policy=ExecutionPolicy.REQUIRES_APPROVAL))

        workflow = Workflow(
            "gated", (Step("w", "writes", build_input=lambda ctx: {"client_id": 7}),)
        )
        engine = WorkflowEngine(registry)
        run = engine.start(workflow)
        run = engine.approve(workflow, run, approved_by="hina@firm.test",
                             amendments={"client_id": 99})
        record = run.step("w")

        # What the AI wanted, what the human changed, and who signed it — the
        # three questions asked when a filing is later disputed.
        assert record.proposed == {"client_id": 7}
        assert record.amendments == {"client_id": 99}
        assert record.approved_by == "hina@firm.test"

    def test_an_approved_step_that_fails_does_not_complete_the_run(self, registry):
        registry.register(
            make_tool(
                "writes",
                policy=ExecutionPolicy.REQUIRES_APPROVAL,
                behaviour=lambda n, kw: ToolResult.failure("CMS rejected it"),
            )
        )

        workflow = Workflow("gated", (Step("w", "writes", max_attempts=1),))
        engine = WorkflowEngine(registry)
        run = engine.approve(workflow, engine.start(workflow))

        # Reporting COMPLETED here would tell a reviewer their approval landed
        # when the document was never filed.
        assert run.state is RunState.FAILED
        assert run.step("w").error == "CMS rejected it"

    def test_a_retry_after_approval_does_not_ask_the_human_again(self, registry):
        registry.register(
            make_tool(
                "writes",
                policy=ExecutionPolicy.REQUIRES_APPROVAL,
                # Fails once — the CMS was briefly unreachable — then succeeds.
                behaviour=lambda n, kw: ToolResult.failure("connection reset")
                if n == 1
                else ToolResult.success({"document_id": 5}),
            )
        )

        workflow = Workflow("gated", (Step("w", "writes", max_attempts=3),))
        engine = WorkflowEngine(registry)
        run = engine.approve(workflow, engine.start(workflow))

        # The approval is held on the step, not spent on one call. Re-pausing
        # would strand the run until someone noticed and signed it twice.
        assert run.state is RunState.COMPLETED
        assert run.context["w"]["document_id"] == 5

    def test_approval_cannot_execute_a_disabled_tool(self, registry):
        """Approval authorises permitted work; it does not grant new powers."""
        tool = make_tool("deletes", policy=ExecutionPolicy.DISABLED)

        result = tool.execute_approved(client_id=7)

        assert not result.ok
        assert tool.calls == []
        assert "approval does not apply" in result.error


class TestAuditTrail:
    def test_every_step_records_its_outcome_and_timing(self, registry):
        registry.register(make_tool("worker"))

        workflow = Workflow("audited", (Step("a", "worker"),))
        run = WorkflowEngine(registry).start(workflow)
        record = run.step("a")

        # "Never lose history": a step nobody can account for afterwards is a
        # step that cannot be promoted to autonomous under ADR-0004.
        assert record.state is StepState.SUCCEEDED
        assert record.attempts == 1
        assert record.started_at is not None
        assert record.finished_at is not None
        assert record.duration_ms is not None

    def test_a_failure_records_its_reason(self, registry):
        registry.register(make_tool("failing", behaviour=lambda n, kw: ToolResult.failure("disk full")))

        workflow = Workflow("audited", (Step("a", "failing", max_attempts=1),))
        run = WorkflowEngine(registry).start(workflow)

        assert run.step("a").error == "disk full"

    def test_confidence_is_carried_onto_the_step(self, registry):
        registry.register(
            make_tool("guesser", behaviour=lambda n, kw: ToolResult.success({"x": 1}, confidence=0.42))
        )

        workflow = Workflow("scored", (Step("a", "guesser"),))
        run = WorkflowEngine(registry).start(workflow)

        # A workflow must be able to branch on how sure a step was.
        assert run.step("a").confidence == 0.42


class TestDefinitionDrift:
    def test_a_stored_run_whose_step_no_longer_exists_fails_loudly(self, registry):
        registry.register(make_tool("first"))

        original = Workflow("v1", (Step("a", "first"), Step("b", "first")))
        engine = WorkflowEngine(registry)
        run = original.new_run()
        engine.store.save(run)

        # The definition changed while a run was in flight.
        renamed = Workflow("v1", (Step("a", "first"),))
        run.step("a").state = StepState.SUCCEEDED
        result = engine.resume(renamed, run)

        # Guessing which step was meant would be worse than stopping.
        assert result.state is RunState.FAILED
        assert "no longer part of" in result.error