import json

import pytest

from agent.orchestration.coordinator import TaskConflict, TaskCoordinator
from agent.orchestration.models import DecisionAction, TaskLedger, TaskStatus, ValidationVerdict
from agent.orchestration.store import (
    FileSystemArtifactStore,
    FileSystemTaskStore,
    TaskStoreError,
)
from agent.trace.store import FileSystemTraceStore

from test_coordinator_integration import (
    FakeExecutor,
    ROOT_TASK_SPEC,
    create_task,
    make_coordinator,
)


@pytest.mark.asyncio
async def test_new_ledger_has_one_stable_root_and_inspect_exposes_it(tmp_path):
    coordinator, store, _ = await make_coordinator(tmp_path, FakeExecutor([]))

    ledger = await store.load("root")
    root_id = ledger.root_task_id
    assert root_id is not None
    assert ledger.focused_task_id == root_id
    assert ledger.tasks[root_id].parent_task_id is None
    assert ledger.tasks[root_id].display_path == "0"
    assert ledger.tasks[root_id].status == TaskStatus.NEEDS_REPLAN
    assert [task.task_id for task in ledger.tasks.values() if task.parent_task_id is None] == [root_id]

    same = await coordinator.ensure_ledger("root", None)
    assert same.root_task_id == root_id
    inspected = await coordinator.inspect_tasks("root")
    assert inspected["root_task_id"] == root_id
    assert inspected["tasks"][0]["current_spec"]["acceptance_criteria"][0][
        "criterion_id"
    ] == "mission-done"


@pytest.mark.asyncio
async def test_root_revise_drives_root_objective_and_mission(tmp_path):
    coordinator, store, _ = await make_coordinator(
        tmp_path, FakeExecutor([])
    )
    ledger = await store.load("root")
    root_id = ledger.root_task_id
    assert ledger.root_objective == "mission"

    await coordinator.decide_task(
        "root",
        root_id,
        None,
        DecisionAction.REVISE,
        {"reason": "clarify", "objective": "revised mission"},
        "revise-root-objective",
    )
    revised = await store.load("root")
    assert revised.root_objective == "revised mission"
    assert revised.mission == "revised mission"


@pytest.mark.asyncio
async def test_rootless_legacy_ledger_falls_back_to_mission(tmp_path):
    store = FileSystemTaskStore(str(tmp_path))
    ledger = TaskLedger(root_trace_id="legacy-rootless", mission="legacy mission")
    assert ledger.root_objective == "legacy mission"
    committed = await store.commit(ledger, expected_revision=-1)
    assert committed.ledger.mission == "legacy mission"
    assert (await store.load("legacy-rootless")).root_objective == "legacy mission"


@pytest.mark.asyncio
async def test_new_ledger_requires_root_spec_and_rootless_policy_is_explicit(tmp_path):
    store = FileSystemTaskStore(str(tmp_path))
    coordinator = TaskCoordinator(
        store,
        FileSystemArtifactStore(str(tmp_path)),
        FileSystemTraceStore(str(tmp_path)),
    )
    with pytest.raises(ValueError, match="root_task_spec is required"):
        await coordinator.ensure_ledger("missing")
    with pytest.raises(TaskConflict, match="has no Root Task ledger"):
        await coordinator.ensure_ledger(
            "missing-resume", ROOT_TASK_SPEC, allow_create=False
        )

    legacy = TaskLedger(root_trace_id="legacy", mission="legacy")
    committed = await store.commit(legacy, expected_revision=-1)
    with pytest.raises(TaskConflict, match="Rootless"):
        await coordinator.ensure_ledger("legacy", ROOT_TASK_SPEC)
    unchanged = await store.load("legacy")
    assert unchanged.root_task_id is None
    assert unchanged.revision == committed.ledger.revision

    invalid_path = store.ledger_path("legacy-invalid")
    invalid_path.parent.mkdir(parents=True, exist_ok=True)
    invalid_path.write_text(json.dumps({
        "root_trace_id": "legacy-invalid",
        "mission": "legacy invalid",
        "tasks": {
            "old-task": {
                "task_id": "old-task",
                "parent_task_id": None,
                "display_path": "1",
                "specs": [{
                    "version": 1,
                    "objective": "old task",
                    "acceptance_criteria": [],
                }],
            }
        },
    }), encoding="utf-8")
    with pytest.raises(TaskStoreError, match="at least one criterion"):
        await store.load("legacy-invalid")


@pytest.mark.asyncio
async def test_parent_waits_until_every_child_is_terminal(tmp_path):
    coordinator, store, _ = await make_coordinator(tmp_path, FakeExecutor([]))
    first = await create_task(coordinator, "first")
    second = await create_task(coordinator, "second")
    third = await create_task(coordinator, "third")
    root_id = (await store.load("root")).root_task_id

    await coordinator.decide_task(
        "root", first, None, DecisionAction.CANCEL, {"reason": "drop first"}, "cancel-first"
    )
    assert (await store.load("root")).tasks[root_id].status == TaskStatus.WAITING_CHILDREN

    superseded = await coordinator.decide_task(
        "root",
        second,
        None,
        DecisionAction.SUPERSEDE,
        {
            "reason": "replace second",
            "replacement": {
                "objective": "second replacement",
                "acceptance_criteria": [
                    {"criterion_id": "replacement", "description": "replacement passes"}
                ],
            },
        },
        "supersede-second",
    )
    replacement = superseded["payload"]["replacement_task_id"]
    assert (await store.load("root")).tasks[root_id].status == TaskStatus.WAITING_CHILDREN

    await coordinator.decide_task(
        "root", replacement, None, DecisionAction.CANCEL, {"reason": "drop replacement"}, "cancel-replacement"
    )
    assert (await store.load("root")).tasks[root_id].status == TaskStatus.WAITING_CHILDREN
    await coordinator.decide_task(
        "root", third, None, DecisionAction.CANCEL, {"reason": "drop third"}, "cancel-third"
    )
    assert (await store.load("root")).tasks[root_id].status == TaskStatus.NEEDS_REPLAN


@pytest.mark.asyncio
async def test_root_executes_after_children_and_receives_only_accepted_results(tmp_path):
    class CapturingExecutor(FakeExecutor):
        def __init__(self):
            super().__init__([ValidationVerdict.PASSED, ValidationVerdict.PASSED])
            self.worker_contexts = []

        async def run_worker(self, context):
            self.worker_contexts.append(context)
            return await super().run_worker(context)

    executor = CapturingExecutor()
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    child_id = await create_task(coordinator, "child")
    root_id = (await store.load("root")).root_task_id

    blocked_dispatch = (await coordinator.dispatch_tasks("root", [root_id]))[0]
    assert "cannot dispatch" in blocked_dispatch.error

    child_cycle = (await coordinator.dispatch_tasks("root", [child_id]))[0]
    await coordinator.decide_task(
        "root",
        child_id,
        child_cycle.validation_id,
        DecisionAction.ACCEPT,
        {"reason": "child accepted"},
        "accept-child",
    )
    accepted_ledger = await store.load("root")
    assert accepted_ledger.tasks[root_id].status == TaskStatus.NEEDS_REPLAN
    assert [
        item["task_id"]
        for item in coordinator._accepted_child_results(
            accepted_ledger, accepted_ledger.tasks[root_id]
        )
    ] == [child_id]

    root_cycle = (await coordinator.dispatch_tasks("root", [root_id]))[0]
    assert root_cycle.error is None
    assert len(executor.worker_contexts) == 2
    child_results = executor.worker_contexts[-1]["accepted_child_results"]
    assert [item["task_id"] for item in child_results] == [child_id]
    assert child_results[0]["attempt_id"] == child_cycle.attempt_id
    assert child_results[0]["validation"]["validation_id"] == child_cycle.validation_id

    await coordinator.decide_task(
        "root",
        root_id,
        root_cycle.validation_id,
        DecisionAction.ACCEPT,
        {"reason": "root accepted"},
        "accept-root",
    )
    completion = await coordinator.mission_completion("root")
    assert completion["status"] == TaskStatus.COMPLETED.value
    assert completion["result_summary"] == "done"


@pytest.mark.asyncio
async def test_root_identity_is_not_cancellable_or_supersedable(tmp_path):
    coordinator, store, _ = await make_coordinator(tmp_path, FakeExecutor([]))
    root_id = (await store.load("root")).root_task_id

    with pytest.raises(TaskConflict, match="root task cannot be cancelled"):
        await coordinator.decide_task(
            "root", root_id, None, DecisionAction.CANCEL, {"reason": "invalid"}, "cancel-root"
        )
    with pytest.raises(TaskConflict, match="root task cannot be superseded"):
        await coordinator.decide_task(
            "root",
            root_id,
            None,
            DecisionAction.SUPERSEDE,
            {
                "reason": "invalid",
                "replacement": {
                    "objective": "replacement",
                    "acceptance_criteria": [
                        {"criterion_id": "replacement", "description": "replacement passes"}
                    ],
                },
            },
            "supersede-root",
        )


@pytest.mark.asyncio
async def test_parent_with_active_descendant_cannot_be_terminated(tmp_path):
    coordinator, store, _ = await make_coordinator(tmp_path, FakeExecutor([]))
    parent_id = await create_task(coordinator, "parent")
    descendant_id = await create_task(
        coordinator, "descendant", parent_task_id=parent_id
    )
    root_id = (await store.load("root")).root_task_id

    with pytest.raises(TaskConflict, match="non-terminal descendants"):
        await coordinator.decide_task(
            "root",
            parent_id,
            None,
            DecisionAction.CANCEL,
            {"reason": "invalid while descendant is active"},
            "cancel-active-parent",
        )
    with pytest.raises(TaskConflict, match="non-terminal descendants"):
        await coordinator.decide_task(
            "root",
            parent_id,
            None,
            DecisionAction.SUPERSEDE,
            {
                "reason": "invalid while descendant is active",
                "replacement": {
                    "objective": "replacement",
                    "acceptance_criteria": [
                        {"criterion_id": "replacement", "description": "replacement passes"}
                    ],
                },
            },
            "supersede-active-parent",
        )
    root_dispatch = (await coordinator.dispatch_tasks("root", [root_id]))[0]
    assert "cannot dispatch" in root_dispatch.error

    await coordinator.decide_task(
        "root",
        descendant_id,
        None,
        DecisionAction.CANCEL,
        {"reason": "finish descendant"},
        "cancel-descendant",
    )
    await coordinator.decide_task(
        "root",
        parent_id,
        None,
        DecisionAction.CANCEL,
        {"reason": "parent can now terminate"},
        "cancel-parent",
    )
    assert (await store.load("root")).tasks[root_id].status == TaskStatus.NEEDS_REPLAN


@pytest.mark.asyncio
async def test_blocked_descendants_allow_mission_to_report_blocked(tmp_path):
    coordinator, store, _ = await make_coordinator(tmp_path, FakeExecutor([]))
    parent_id = await create_task(coordinator, "parent")
    child_id = await create_task(coordinator, "child", parent_task_id=parent_id)
    root_id = (await store.load("root")).root_task_id

    await coordinator.decide_task(
        "root",
        child_id,
        None,
        DecisionAction.BLOCK,
        {"reason": "child external dependency"},
        "block-child",
    )
    await coordinator.decide_task(
        "root",
        parent_id,
        None,
        DecisionAction.BLOCK,
        {"reason": "parent waits on blocked child"},
        "block-parent",
    )
    await coordinator.decide_task(
        "root",
        root_id,
        None,
        DecisionAction.BLOCK,
        {"reason": "mission external dependency"},
        "block-root",
    )

    completion = await coordinator.mission_completion("root")
    assert completion["status"] == TaskStatus.BLOCKED.value
    assert completion["blocked_reason"] == "mission external dependency"

    await coordinator.decide_task(
        "root",
        root_id,
        None,
        DecisionAction.UNBLOCK,
        {},
        "unblock-root",
    )
    assert (await store.load("root")).tasks[root_id].status == TaskStatus.NEEDS_REPLAN


@pytest.mark.asyncio
async def test_parent_cannot_block_while_descendant_attempt_is_active(tmp_path):
    coordinator, _, _ = await make_coordinator(tmp_path, FakeExecutor([]))
    parent_id = await create_task(coordinator, "parent")
    child_id = await create_task(coordinator, "child", parent_task_id=parent_id)
    await coordinator._create_attempt("root", child_id, "worker", None)

    with pytest.raises(TaskConflict, match="active descendants"):
        await coordinator.decide_task(
            "root",
            parent_id,
            None,
            DecisionAction.BLOCK,
            {"reason": "must not hide active work"},
            "block-active-parent",
        )


@pytest.mark.asyncio
async def test_inserting_a_sibling_rebases_its_existing_subtree(tmp_path):
    coordinator, store, _ = await make_coordinator(tmp_path, FakeExecutor([]))
    first_id = await create_task(coordinator, "first")
    moved_id = await create_task(coordinator, "moved")
    descendant_id = await create_task(
        coordinator, "moved descendant", parent_task_id=moved_id
    )

    await coordinator.create_tasks(
        "root",
        [{
            "objective": "inserted",
            "acceptance_criteria": [
                {"criterion_id": "inserted", "description": "inserted passes"}
            ],
        }],
        placement={"after_task_id": first_id},
    )
    ledger = await store.load("root")
    assert ledger.tasks[moved_id].display_path == "0.3"
    assert ledger.tasks[descendant_id].display_path == "0.3.1"


@pytest.mark.asyncio
async def test_invalid_batch_and_revision_are_atomic(tmp_path):
    executor = FakeExecutor([ValidationVerdict.FAILED])
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    root_id = (await store.load("root")).root_task_id

    with pytest.raises(ValueError, match="at least one criterion"):
        await coordinator.create_tasks(
            "root",
            [
                {
                    "objective": "valid first draft",
                    "acceptance_criteria": [
                        {"criterion_id": "valid", "description": "valid passes"}
                    ],
                },
                {"objective": "invalid second draft"},
            ],
        )
    ledger = await store.load("root")
    assert list(ledger.tasks) == [root_id]
    assert ledger.tasks[root_id].status == TaskStatus.NEEDS_REPLAN

    with pytest.raises(ValueError, match="context_refs"):
        await coordinator.create_tasks(
            "root",
            [{
                "objective": "invalid context refs",
                "acceptance_criteria": [
                    {"criterion_id": "context", "description": "context passes"}
                ],
                "context_refs": "memory://not-a-list",
            }],
        )

    task_id = await create_task(coordinator, "revise atomically")
    cycle = (await coordinator.dispatch_tasks("root", [task_id]))[0]
    for index, criteria in enumerate(([], None, "criterion")):
        with pytest.raises(ValueError, match="acceptance_criteria|at least one criterion"):
            await coordinator.decide_task(
                "root",
                task_id,
                cycle.validation_id,
                DecisionAction.REVISE,
                {"reason": "invalid revision", "acceptance_criteria": criteria},
                f"invalid-revision-{index}",
            )
    task = (await store.load("root")).tasks[task_id]
    assert task.current_spec_version == 1
    assert len(task.specs) == 1
    assert task.status == TaskStatus.AWAITING_DECISION

    for index, objective in enumerate((123, None, "", "   ")):
        with pytest.raises(ValueError, match="objective"):
            await coordinator.decide_task(
                "root",
                task_id,
                cycle.validation_id,
                DecisionAction.REVISE,
                {"reason": "invalid objective", "objective": objective},
                f"invalid-objective-{index}",
            )
        unchanged = (await store.load("root")).tasks[task_id]
        assert unchanged.current_spec_version == 1
        assert unchanged.status == TaskStatus.AWAITING_DECISION