from pathlib import Path

import pytest

from src.events import record_event
from src.executor import (
    WorkflowResumeError,
    execute_workflow,
    resume_workflow,
)
from src.persistence import (
    InMemoryStateStore,
    RunNotFoundError,
    SQLiteStateStore,
)
from src.registry import HandlerRegistry
from src.schemas import (
    EventType,
    NodeDefinition,
    NodeRun,
    NodeStatus,
    NodeType,
    WorkflowDefinition,
    WorkflowRun,
    WorkflowStatus,
)


def make_two_step_workflow(*, version: str = "1.0") -> WorkflowDefinition:
    return WorkflowDefinition(
        workflow_id="persistence-workflow",
        name="Persistence Workflow",
        version=version,
        nodes=[
            NodeDefinition(
                node_id="first",
                name="First",
                node_type=NodeType.TASK,
                handler="first_handler",
            ),
            NodeDefinition(
                node_id="second",
                name="Second",
                node_type=NodeType.TASK,
                handler="second_handler",
                depends_on=["first"],
            ),
        ],
    )


def test_in_memory_store_returns_detached_snapshot() -> None:
    store = InMemoryStateStore()
    run = WorkflowRun(
        run_id="run-1",
        workflow_id="workflow-1",
        context={"value": 1},
    )

    store.save(run)
    run.context["value"] = 99

    loaded = store.load("run-1")

    assert loaded.context == {"value": 1}
    loaded.context["value"] = 7
    assert store.load("run-1").context == {"value": 1}


def test_in_memory_store_rejects_unknown_run() -> None:
    store = InMemoryStateStore()

    with pytest.raises(RunNotFoundError, match="was not found"):
        store.load("missing")


def test_sqlite_store_round_trips_nested_workflow_state(tmp_path: Path) -> None:
    store = SQLiteStateStore(tmp_path / "workflow.db")
    run = WorkflowRun(
        run_id="sqlite-run",
        workflow_id="workflow-1",
        workflow_version="2.1",
        status=WorkflowStatus.RUNNING,
        node_runs={
            "task": NodeRun(
                node_id="task",
                status=NodeStatus.COMPLETED,
                attempt=1,
                output={"ok": True},
            )
        },
        context={"request": "synthetic"},
    )
    record_event(run, EventType.WORKFLOW_STARTED)
    record_event(run, EventType.NODE_COMPLETED, node_id="task")

    store.save(run)
    loaded = store.load("sqlite-run")

    assert loaded == run
    assert loaded.events[0].event_type is EventType.WORKFLOW_STARTED
    assert loaded.node_runs["task"].output == {"ok": True}


def test_sqlite_store_upserts_latest_checkpoint(tmp_path: Path) -> None:
    store = SQLiteStateStore(tmp_path / "workflow.db")
    run = WorkflowRun(run_id="run-1", workflow_id="workflow-1")

    store.save(run)
    run.status = WorkflowStatus.COMPLETED
    run.final_output = {"done": True}
    store.save(run)

    loaded = store.load("run-1")

    assert loaded.status is WorkflowStatus.COMPLETED
    assert loaded.final_output == {"done": True}


def test_sqlite_store_rejects_unknown_run(tmp_path: Path) -> None:
    store = SQLiteStateStore(tmp_path / "workflow.db")

    with pytest.raises(RunNotFoundError, match="was not found"):
        store.load("missing")


class SimulatedProcessStop(BaseException):
    pass


class StopAfterFirstCompletionStore:
    """Save a safe checkpoint, then simulate process termination once."""

    def __init__(self) -> None:
        self.inner = InMemoryStateStore()
        self.triggered = False

    def save(self, run: WorkflowRun) -> None:
        self.inner.save(run)
        first = run.node_runs.get("first")
        second = run.node_runs.get("second")
        if (
            not self.triggered
            and first is not None
            and second is not None
            and first.status is NodeStatus.COMPLETED
            and second.status is NodeStatus.PENDING
        ):
            self.triggered = True
            raise SimulatedProcessStop()

    def load(self, run_id: str) -> WorkflowRun:
        return self.inner.load(run_id)


def test_safe_checkpoint_resume_skips_completed_node_after_process_stop() -> None:
    registry = HandlerRegistry()
    first_calls = 0
    second_calls = 0

    def first_handler(payload: dict) -> dict:
        nonlocal first_calls
        first_calls += 1
        return {"first": "done"}

    def second_handler(payload: dict) -> dict:
        nonlocal second_calls
        second_calls += 1
        assert payload["dependencies"] == {"first": {"first": "done"}}
        return {"second": "done"}

    registry.register("first_handler", first_handler)
    registry.register("second_handler", second_handler)
    workflow = make_two_step_workflow()
    store = StopAfterFirstCompletionStore()

    with pytest.raises(SimulatedProcessStop):
        execute_workflow(
            workflow,
            registry,
            run_id="resume-run",
            state_store=store,
        )

    checkpoint = store.load("resume-run")
    assert checkpoint.status is WorkflowStatus.RUNNING
    assert checkpoint.node_runs["first"].status is NodeStatus.COMPLETED
    assert checkpoint.node_runs["second"].status is NodeStatus.PENDING
    assert first_calls == 1
    assert second_calls == 0

    resumed = resume_workflow(
        workflow,
        registry,
        run_id="resume-run",
        state_store=store,
    )

    assert resumed.status is WorkflowStatus.COMPLETED
    assert resumed.node_runs["first"].attempt == 1
    assert resumed.node_runs["second"].attempt == 1
    assert first_calls == 1
    assert second_calls == 1
    assert EventType.WORKFLOW_RESUMED in [event.event_type for event in resumed.events]
    assert resumed.final_output == {"second": {"second": "done"}}


def test_completed_run_resume_is_idempotent() -> None:
    registry = HandlerRegistry()
    calls = 0

    def handler(payload: dict) -> dict:
        nonlocal calls
        calls += 1
        return {"ok": True}

    registry.register("first_handler", handler)
    registry.register("second_handler", handler)
    workflow = make_two_step_workflow()
    store = InMemoryStateStore()

    original = execute_workflow(
        workflow,
        registry,
        run_id="completed-run",
        state_store=store,
    )
    original_event_ids = [event.event_id for event in original.events]

    resumed = resume_workflow(
        workflow,
        registry,
        run_id="completed-run",
        state_store=store,
    )

    assert calls == 2
    assert resumed.status is WorkflowStatus.COMPLETED
    assert [event.event_id for event in resumed.events] == original_event_ids


def test_resume_rejects_changed_workflow_version() -> None:
    store = InMemoryStateStore()
    run = WorkflowRun(
        run_id="run-1",
        workflow_id="persistence-workflow",
        workflow_version="1.0",
        status=WorkflowStatus.RUNNING,
        node_runs={
            "first": NodeRun(node_id="first", status=NodeStatus.COMPLETED),
            "second": NodeRun(node_id="second", status=NodeStatus.PENDING),
        },
    )
    store.save(run)

    with pytest.raises(WorkflowResumeError, match="version"):
        resume_workflow(
            make_two_step_workflow(version="2.0"),
            HandlerRegistry(),
            run_id="run-1",
            state_store=store,
        )


def test_resume_rejects_ambiguous_in_flight_node() -> None:
    store = InMemoryStateStore()
    run = WorkflowRun(
        run_id="run-1",
        workflow_id="persistence-workflow",
        workflow_version="1.0",
        status=WorkflowStatus.RUNNING,
        node_runs={
            "first": NodeRun(node_id="first", status=NodeStatus.COMPLETED),
            "second": NodeRun(node_id="second", status=NodeStatus.RUNNING, attempt=1),
        },
    )
    store.save(run)

    with pytest.raises(WorkflowResumeError, match="ambiguous or non-resumable"):
        resume_workflow(
            make_two_step_workflow(),
            HandlerRegistry(),
            run_id="run-1",
            state_store=store,
        )