from __future__ import annotations

import json
from dataclasses import replace
from datetime import UTC, datetime
from types import SimpleNamespace

import pytest
from agent import AgentRunner, FileSystemArtifactStore, FileSystemTaskStore, FileSystemTraceStore
from agent.orchestration import ArtifactRef, TaskStatus, ValidationVerdict

from script_build_host.composition import HostDependencies, compose_host
from script_build_host.domain.artifacts import (
    ArtifactKind,
    ArtifactState,
    ArtifactVersion,
    EvidenceRecordV1,
)
from script_build_host.domain.errors import ArtifactNotFound
from script_build_host.domain.input_snapshot import ScriptBuildInputSnapshotV1
from script_build_host.domain.records import (
    BuildStatus,
    MissionBinding,
    Principal,
    PublicationState,
)
from script_build_host.infrastructure.canonical_json import canonical_sha256
from script_build_host.tools.gateway import RetrievalResult


def _tool(call_id: str, name: str, arguments: dict[str, object]) -> dict[str, object]:
    return {
        "content": "",
        "tool_calls": [
            {
                "id": call_id,
                "type": "function",
                "function": {"name": name, "arguments": json.dumps(arguments)},
            }
        ],
        "finish_reason": "tool_calls",
    }


def _prompt(messages: list[dict[str, object]]) -> dict[str, object]:
    for message in reversed(messages):
        if message.get("role") != "user" or not isinstance(message.get("content"), str):
            continue
        try:
            value = json.loads(message["content"])
        except ValueError:
            continue
        if isinstance(value, dict) and "task_spec" in value:
            return value
    raise AssertionError("missing role prompt")


def _last_tool_json(messages: list[dict[str, object]]) -> dict[str, object] | None:
    for message in reversed(messages):
        if message.get("role") != "tool":
            continue
        content = message.get("content")
        if isinstance(content, str):
            try:
                value = json.loads(content)
                return value if isinstance(value, dict) else None
            except ValueError:
                return None
    return None


def _contract(
    kind: str,
    objective: str,
    criterion_id: str,
    description: str,
    *,
    scope: str,
) -> dict[str, object]:
    return {
        "task_kind": kind,
        "scope_ref": scope,
        "intent_class": "explore",
        "objective": objective,
        "input_decision_ids": [],
        "base_decision_id": None,
        "gap_ref": None,
        "criteria": [
            {
                "criterion_id": criterion_id,
                "description": description,
                "hard": True,
            }
        ],
        "goal_ids": [],
        "supersedes_decision_ids": [],
        "comparison_decision_ids": [],
    }


class _ScriptedLLM:
    def __init__(self, store: FileSystemTaskStore, root: str) -> None:
        self.store = store
        self.root = root
        self.calls = 0
        self.decode_versions: list[int] = []
        self.validated_artifact_kinds: list[str] = []
        self.replayed_dispatch = False

    def call(self, name: str, arguments: dict[str, object]) -> dict[str, object]:
        self.calls += 1
        return _tool(f"call-{self.calls}", name, arguments)

    def call_with_id(
        self,
        call_id: str,
        name: str,
        arguments: dict[str, object],
    ) -> dict[str, object]:
        self.calls += 1
        return _tool(call_id, name, arguments)

    async def __call__(self, messages, tools, **_kwargs):
        names = {item["function"]["name"] for item in tools or []}
        if "plan_script_tasks" in names:
            return await self._planner()
        if "submit_validation" in names:
            return self._validator(messages)
        if "submit_attempt" in names:
            return self._worker(messages, names)
        raise AssertionError(sorted(names))

    async def _planner(self):
        ledger = await self.store.load(self.root)
        root = ledger.tasks[ledger.root_task_id]
        direction = next(
            (
                item
                for item in ledger.tasks.values()
                if "script-build://task-kinds/direction" in item.current_spec.context_refs
            ),
            None,
        )
        if direction is None:
            return self.call(
                "plan_script_tasks",
                {
                    "contracts": [
                        _contract(
                            "direction",
                            "build direction",
                            "direction-ready",
                            "direction is grounded",
                            scope="script-build://scopes/direction",
                        )
                    ],
                },
            )
        retrieval = next(
            (
                item
                for item in ledger.tasks.values()
                if "script-build://task-kinds/decode-retrieval" in item.current_spec.context_refs
            ),
            None,
        )
        if retrieval is None:
            return self.call(
                "plan_script_tasks",
                {
                    "parent_task_id": direction.task_id,
                    "contracts": [
                        _contract(
                            "decode-retrieval",
                            "decode evidence",
                            "evidence-ready",
                            "evidence is relevant",
                            scope="script-build://scopes/direction/decode",
                        )
                    ],
                },
            )
        if retrieval.status in {TaskStatus.PENDING, TaskStatus.NEEDS_REPLAN}:
            return self.call_with_id(
                f"decode-dispatch-v{retrieval.current_spec_version}",
                "dispatch_script_tasks",
                {"task_ids": [retrieval.task_id]},
            )
        if retrieval.status == TaskStatus.AWAITING_DECISION:
            if retrieval.current_spec_version == 1 and not self.replayed_dispatch:
                self.replayed_dispatch = True
                return self.call_with_id(
                    "decode-dispatch-v1",
                    "dispatch_script_tasks",
                    {"task_ids": [retrieval.task_id]},
                )
            validation = ledger.validations[retrieval.validation_ids[-1]]
            if validation.verdict == ValidationVerdict.FAILED:
                return self.call(
                    "decide_script_task",
                    {
                        "task_id": retrieval.task_id,
                        "validation_id": validation.validation_id,
                        "action": "revise",
                        "reason": "make decode query precise",
                        "replacement_contract": _contract(
                            "decode-retrieval",
                            "decode evidence revised",
                            "evidence-ready",
                            "evidence is relevant",
                            scope="script-build://scopes/direction/decode",
                        ),
                    },
                )
            return self.call(
                "decide_script_task",
                {
                    "task_id": retrieval.task_id,
                    "validation_id": validation.validation_id,
                    "action": "accept",
                    "reason": "decode evidence passed",
                },
            )
        if direction.status in {TaskStatus.PENDING, TaskStatus.NEEDS_REPLAN}:
            return self.call("dispatch_script_tasks", {"task_ids": [direction.task_id]})
        if direction.status == TaskStatus.AWAITING_DECISION:
            validation = ledger.validations[direction.validation_ids[-1]]
            return self.call(
                "decide_script_task",
                {
                    "task_id": direction.task_id,
                    "validation_id": validation.validation_id,
                    "action": "accept",
                    "reason": "direction passed",
                },
            )
        assert direction.status == TaskStatus.COMPLETED
        if root.current_spec_version == 1:
            return self.call(
                "decide_script_task",
                {
                    "task_id": root.task_id,
                    "action": "revise",
                    "reason": "record accepted direction in root objective",
                    "replacement_contract": {
                        "objective": root.current_spec.objective + " using accepted direction",
                        "acceptance_criteria": [
                            {
                                "criterion_id": item.criterion_id,
                                "description": item.description,
                                "hard": item.hard,
                            }
                            for item in root.current_spec.acceptance_criteria
                        ],
                        "context_refs": list(root.current_spec.context_refs),
                    },
                },
            )
        return self.call(
            "decide_script_task",
            {
                "task_id": root.task_id,
                "action": "block",
                "reason": "PHASE_ONE_CAPABILITY_BOUNDARY",
            },
        )

    def _worker(self, messages, names):
        prompt = _prompt(messages)
        spec = prompt["task_spec"]
        last = _last_tool_json(messages)
        if "search_script_decode_case" in names:
            if last is None or "artifact_ref" not in last:
                self.decode_versions.append(spec["version"])
                return self.call(
                    "search_script_decode_case",
                    {"return_field": "主脉络", "keyword": "focus", "top_k": 3},
                )
            return self.call(
                "submit_attempt",
                {},
            )
        if last is None or "artifact_ref" not in last:
            accepted = prompt["accepted_child_results"]
            evidence_uri = accepted[0]["submission"]["evidence_refs"][0]["uri"]
            return self.call(
                "save_direction_candidate",
                {
                    "goals": [
                        {
                            "client_key": "grounded-direction",
                            "statement": "make one grounded direction",
                            "rationale": "accepted evidence",
                            "success_criteria": ["the direction remains grounded"],
                        }
                    ],
                    "constraints": [
                        {
                            "constraint_id": "c1",
                            "statement": "remain grounded",
                            "rationale": "accepted evidence is mandatory",
                        }
                    ],
                    "preferences": [],
                    "evidence_refs": [evidence_uri],
                },
            )
        return self.call(
            "submit_attempt",
            {},
        )

    def _validator(self, messages):
        prompt = _prompt(messages)
        spec = prompt["task_spec"]
        kind = (
            "direction"
            if "script-build://task-kinds/direction" in spec["context_refs"]
            else "evidence"
        )
        evidence = _last_tool_json(messages)
        if evidence is None or "items" not in evidence:
            return self.call(
                "query_validation_evidence",
                {"query": kind, "limit": 5},
            )
        assert evidence["items"], "validator must read the frozen artifact body"
        artifact = evidence["items"][0]["artifact"]
        assert isinstance(artifact, dict)
        self.validated_artifact_kinds.append(kind)
        failed = spec["objective"] == "decode evidence" and spec["version"] == 1
        verdict = "failed" if failed else "passed"
        snapshot = prompt["artifact_snapshot"]
        refs = [*snapshot["artifact_refs"], *snapshot["evidence_refs"]]
        artifact_ref = refs[0]
        return self.call(
            "submit_validation",
            {
                "verdict": verdict,
                "criterion_results": [
                    {
                        "criterion_id": item["criterion_id"],
                        "verdict": verdict,
                        "reason": "checked frozen snapshot",
                        "evidence_refs": [] if failed else [artifact_ref],
                    }
                    for item in spec["acceptance_criteria"]
                ],
                "defects": (
                    [
                        {
                            "defect_code": "DECODE_QUERY_TOO_BROAD",
                            "criterion_id": spec["acceptance_criteria"][0]["criterion_id"],
                            "scope_ref": "script-build://scopes/direction/decode",
                            "observed_excerpt": "v1 decode evidence was too broad",
                            "evidence_refs": [],
                            "severity": "hard",
                            "invalidated_inputs": [],
                            "recommended_action_class": "revise",
                        }
                    ]
                    if failed
                    else []
                ),
                "recommendation": "revise" if failed else "accept",
            },
        )


class _Bindings:
    def __init__(self, binding: MissionBinding) -> None:
        self.binding = binding
        self.active = None

    async def get_by_build(self, _: int) -> MissionBinding:
        return self.binding

    async def get_by_root(self, _: str) -> MissionBinding:
        return self.binding

    async def set_active_direction(self, **values) -> MissionBinding:
        self.active = values["artifact_version_id"]
        return replace(self.binding, active_direction_artifact_version_id=self.active)


class _Snapshots:
    def __init__(self, snapshot: ScriptBuildInputSnapshotV1) -> None:
        self.snapshot = snapshot

    async def get(self, *_args, **_kwargs) -> ScriptBuildInputSnapshotV1:
        return self.snapshot


class _Artifacts:
    def __init__(self) -> None:
        self.values = {}
        self.next_id = 1

    async def freeze(self, *, artifact, script_build_id, task_id, attempt_id, spec_version):
        payload = artifact.content_payload()
        digest = canonical_sha256(payload).wire
        if isinstance(artifact, EvidenceRecordV1):
            artifact = replace(artifact, content_sha256=digest)
            kind = ArtifactKind.EVIDENCE
        else:
            artifact = replace(artifact, canonical_sha256=digest)
            kind = ArtifactKind.DIRECTION
        value = ArtifactVersion(
            self.next_id,
            script_build_id,
            task_id,
            attempt_id,
            spec_version,
            kind,
            digest,
            ArtifactState.FROZEN,
            artifact,
            datetime.now(UTC),
            datetime.now(UTC),
        )
        ref = ArtifactRef(
            uri=f"script-build://artifact-versions/{self.next_id}",
            kind=kind.value,
            version=str(self.next_id),
            digest=digest,
            metadata={
                "script_build_id": script_build_id,
                "task_id": task_id,
                "attempt_id": attempt_id,
            },
        )
        self.values[self.next_id] = (value, ref)
        self.next_id += 1
        return value, ref

    async def read_by_ref(self, ref, **_kwargs):
        return self.values[int(ref.version)][0]

    async def get_by_attempt(self, *, script_build_id, task_id, attempt_id):
        try:
            return next(
                value
                for value, _ref in self.values.values()
                if value.script_build_id == script_build_id
                and value.task_id == task_id
                and value.attempt_id == attempt_id
            )
        except StopIteration as exc:
            raise ArtifactNotFound() from exc


class _Legacy:
    def __init__(self) -> None:
        self.status = BuildStatus.RUNNING
        self.direction = None

    async def get_status(self, _: int) -> BuildStatus:
        return self.status

    async def set_status(self, _id, status, **_kwargs):
        self.status = status

    async def project_direction(self, _id, value):
        assert self.status != BuildStatus.STOPPING
        self.direction = value


class _Publications:
    def __init__(self) -> None:
        self.state = PublicationState.PENDING

    async def prepare(self, **values):
        return SimpleNamespace(publication_id=1, state=self.state, **values)

    async def mark_published(self, _id):
        self.state = PublicationState.PUBLISHED

    async def mark_failed(self, _id, **_values):
        self.state = PublicationState.FAILED


class _Decode:
    async def retrieve(self, *, query, snapshot):
        del snapshot
        return RetrievalResult(
            source_refs=(f"decode:{query['keyword']}",),
            summary="evidence",
            supports=("direction",),
            confidence="high",
            metadata={"index_sha256": "sha256:" + "1" * 64, "ranks": [1]},
        )


class _PrincipalProvider:
    async def current(self, request_context=None):
        del request_context
        return Principal("owner")


class _Authorizer:
    async def require_source_access(self, _principal, **_source):
        return None

    async def require_access(self, _principal, _build):
        return None


@pytest.mark.asyncio
async def test_real_runner_phase_one_fail_revise_pass_direction_accept_then_root_block(tmp_path):
    now = datetime.now(UTC)
    binding = MissionBinding(1, 7, "phase-one-root", 11, None, None, "test", "v1", now, now)
    model_presets = {
        name: {"model": "fake-model", "temperature": 0, "max_iterations": 50}
        for name in (
            "script_decode_retrieval_worker",
            "script_direction_worker",
            "script_retrieval_validator",
            "script_candidate_validator",
        )
    }
    snapshot = ScriptBuildInputSnapshotV1(
        "11",
        7,
        1,
        2,
        3,
        {"topic": {"id": 3, "result": "topic"}},
        {"account_name": "acct"},
        (),
        (),
        (),
        (),
        {},
        {"presets": model_presets},
        "sha256:" + "a" * 64,
        now,
    )
    task_store = FileSystemTaskStore(str(tmp_path / "ledger"))
    llm = _ScriptedLLM(task_store, binding.root_trace_id)
    runner = AgentRunner(trace_store=FileSystemTraceStore(str(tmp_path / "trace")), llm_call=llm)
    bindings = _Bindings(binding)
    snapshots = _Snapshots(snapshot)
    artifacts = _Artifacts()
    legacy = _Legacy()
    publications = _Publications()
    composition = compose_host(
        HostDependencies(
            runner=runner,
            task_store=task_store,
            framework_artifact_store=FileSystemArtifactStore(str(tmp_path / "artifacts")),
            input_snapshot_service=SimpleNamespace(get=snapshots.get),
            input_snapshots=snapshots,
            bindings=bindings,
            business_artifacts=artifacts,
            publications=publications,
            legacy_state=legacy,
            principal_provider=_PrincipalProvider(),
            build_authorizer=_Authorizer(),
            retrieval_adapters={"decode": _Decode()},
        )
    )
    await composition.mission_service.run(7)
    ledger = await task_store.load(binding.root_trace_id)
    root = ledger.tasks[ledger.root_task_id]
    assert root.status == TaskStatus.BLOCKED
    assert root.blocked_reason == "PHASE_ONE_CAPABILITY_BOUNDARY"
    assert root.current_spec_version == 2
    assert llm.decode_versions == [1, 2]
    assert llm.replayed_dispatch
    assert len(ledger.operations) == 3
    assert llm.validated_artifact_kinds == ["evidence", "evidence", "direction"]
    assert legacy.status == BuildStatus.PARTIAL
    assert legacy.direction.startswith("# Direction\n\n## Goals")
    assert "### goal-" in legacy.direction
    assert ": make one grounded direction" in legacy.direction
    assert "## Constraints" in legacy.direction
    assert publications.state == PublicationState.PUBLISHED
    assert bindings.active is not None
    assert not hasattr(bindings.binding, "branch_id")

    model_calls = llm.calls
    reloaded_ledger = await FileSystemTaskStore(str(tmp_path / "ledger")).load(
        binding.root_trace_id
    )
    assert reloaded_ledger.tasks[reloaded_ledger.root_task_id].status == TaskStatus.BLOCKED
    reloaded_traces = FileSystemTraceStore(str(tmp_path / "trace"))
    reloaded_trace = await reloaded_traces.get_trace(binding.root_trace_id)
    assert reloaded_trace is not None
    assert await reloaded_traces.get_goal_tree(binding.root_trace_id) is None
    submitted = next(
        attempt for attempt in reloaded_ledger.attempts.values() if attempt.snapshot_id
    )
    assert await FileSystemArtifactStore(str(tmp_path / "artifacts")).get(
        binding.root_trace_id,
        submitted.snapshot_id,
    )
    assert llm.calls == model_calls