from __future__ import annotations

from dataclasses import replace
from datetime import UTC, datetime
from hashlib import sha256
from types import SimpleNamespace

import pytest
from agent import AgentRole, ToolRegistry, get_preset

from script_build_host.agents.model_resolver import (
    ScriptRoleRunConfigResolver,
    SnapshotModelManifestSource,
)
from script_build_host.agents.presets import register_script_presets
from script_build_host.agents.prompt_resolver import SnapshotRoleSystemPromptResolver
from script_build_host.agents.prompts import (
    phase_three_prompt_manifest,
    script_build_prompt_manifest,
)
from script_build_host.domain.canonical_json import canonical_sha256
from script_build_host.domain.errors import ProtocolViolation
from script_build_host.domain.input_snapshot import ScriptBuildInputSnapshotV1
from script_build_host.tools.gateway import LegacyScriptToolGateway
from script_build_host.tools.registry import register_script_tools


def test_script_presets_are_exact_and_never_expose_legacy_control_tools() -> None:
    register_script_presets()
    expected = {
        "script_planner",
        "script_direction_worker",
        "script_pattern_retrieval_worker",
        "script_decode_retrieval_worker",
        "script_external_retrieval_worker",
        "script_knowledge_retrieval_worker",
        "script_retrieval_validator",
        "script_candidate_validator",
        "script_structure_worker",
        "script_paragraph_worker",
        "script_element_set_worker",
        "script_candidate_compare_worker",
        "script_compose_worker",
        "script_candidate_portfolio_worker",
    }
    forbidden = {
        "agent",
        "evaluate",
        "bash_command",
        "write_file",
        "dispatch_tasks",
        "implement_paths",
        "begin_round",
        "multipath",
    }
    for name in expected:
        preset = get_preset(name)
        assert preset.role in {AgentRole.PLANNER, AgentRole.WORKER, AgentRole.VALIDATOR}
        assert forbidden.isdisjoint(set(preset.allowed_tools or []))
        assert set(preset.allowed_tools or []).isdisjoint(set(preset.denied_tools or []))

    planner = get_preset("script_planner")
    assert set(planner.allowed_tools or []) == {
        "plan_script_tasks",
        "decide_script_task",
        "inspect_script_plan",
        "dispatch_script_tasks",
        "validate_attempt",
        "read_input_snapshot",
    }
    compose = get_preset("script_compose_worker")
    assert "read_active_frontier" in (compose.allowed_tools or [])
    assert "read_accepted_artifact" not in (compose.allowed_tools or [])
    compare = get_preset("script_candidate_compare_worker")
    assert "read_pinned_candidates" in (compare.allowed_tools or [])
    assert "create_script_paragraph" not in (compare.allowed_tools or [])
    portfolio = get_preset("script_candidate_portfolio_worker")
    assert "save_candidate_portfolio" in (portfolio.allowed_tools or [])
    assert "save_structured_script_candidate" not in (portfolio.allowed_tools or [])
    for validator_name in ("script_retrieval_validator", "script_candidate_validator"):
        assert "view_frozen_images" in (get_preset(validator_name).allowed_tools or [])

    assert set(get_preset("script_root_worker").allowed_tools or []) == {
        "read_input_snapshot",
        "read_accepted_portfolio",
        "legacy_projection_dry_run",
        "save_root_delivery_manifest",
        "submit_attempt",
    }
    assert set(get_preset("script_root_validator").allowed_tools or []) == {
        "read_input_snapshot",
        "read_accepted_portfolio",
        "query_validation_evidence",
        "deterministic_precheck",
        "submit_validation",
    }


def test_new_build_prompt_manifest_freezes_every_phase_one_and_phase_two_role() -> None:
    manifest = script_build_prompt_manifest()
    assert len(manifest) == 14
    assert len({str(item["preset"]) for item in manifest}) == len(manifest)
    assert {str(item["preset"]) for item in manifest} >= {
        "script_planner",
        "script_structure_worker",
        "script_paragraph_worker",
        "script_element_set_worker",
        "script_candidate_compare_worker",
        "script_compose_worker",
        "script_candidate_portfolio_worker",
    }
    for item in manifest:
        assert str(item["content"])
        assert str(item["content_sha256"]).startswith("sha256:")
        assert len(str(item["content_sha256"])) == 71


def test_phase_three_prompt_manifest_is_versioned_separately() -> None:
    manifest = phase_three_prompt_manifest()
    assert {str(item["preset"]) for item in manifest} == {
        "script_root_worker",
        "script_root_validator",
    }
    assert all(str(item["content_sha256"]).startswith("sha256:") for item in manifest)


class _Gateway:
    coordinator = SimpleNamespace(task_store=None)


def test_legacy_retrieval_tool_schemas_keep_old_argument_shapes() -> None:
    registry = ToolRegistry()
    register_script_tools(registry, _Gateway())  # type: ignore[arg-type]
    schemas = {
        item["function"]["name"]: item["function"]["parameters"] for item in registry.get_schemas()
    }
    decode = schemas["search_script_decode_case"]
    assert decode["required"] == ["return_field"]
    assert decode["properties"]["match_fields"]["type"] == "object"
    assert decode["properties"]["top_k"]["default"] == 3
    external = schemas["external_search_case"]
    assert external["required"] == ["keyword"]
    assert external["properties"]["platform_channel"]["default"] == "xhs"
    knowledge = schemas["search_knowledge"]
    assert knowledge["required"] == ["keyword"]
    assert knowledge["properties"]["max_count"]["default"] == 3
    assert schemas["query_pattern_qa"]["required"] == ["message"]
    assert schemas["load_images"]["required"] == ["image_urls"]
    assert schemas["load_frozen_strategy"]["required"] == ["strategy_ref"]
    assert schemas["submit_attempt"].get("properties") == {}
    assert schemas["submit_validation"]["required"] == [
        "verdict",
        "criterion_results",
        "defects",
    ]
    assert "summary" not in schemas["submit_attempt"].get("properties", {})
    assert "artifact_refs" not in schemas["submit_attempt"].get("properties", {})
    assert {item.value for item in registry.get_capabilities("search_script_decode_case")} == {
        "external_send",
        "read",
        "write",
    }


class _Bindings:
    async def get_by_root(self, root_trace_id: str) -> SimpleNamespace:
        assert root_trace_id == "root-a"
        return SimpleNamespace(script_build_id=9, input_snapshot_id=3)


class _Snapshots:
    async def get(self, snapshot_id: str, *, script_build_id: int) -> ScriptBuildInputSnapshotV1:
        assert (snapshot_id, script_build_id) == ("3", 9)
        return ScriptBuildInputSnapshotV1(
            snapshot_id="3",
            script_build_id=9,
            execution_id=1,
            topic_build_id=2,
            topic_id=3,
            topic={},
            account={},
            persona_points=(),
            section_patterns=(),
            strategies=(),
            prompt_manifest=(),
            datasource_manifest={},
            model_manifest={
                "presets": {
                    "script_direction_worker": {
                        "model": "frozen-direction-model",
                        "temperature": 0.2,
                        "max_iterations": 18,
                    }
                }
            },
            canonical_sha256="sha256:" + "a" * 64,
            created_at=datetime.now(UTC),
        )


@pytest.mark.asyncio
async def test_role_model_resolver_reads_each_roots_frozen_manifest() -> None:
    resolver = ScriptRoleRunConfigResolver(
        {}, manifest_source=SnapshotModelManifestSource(_Bindings(), _Snapshots())
    )
    value = await resolver.resolve(
        role=AgentRole.WORKER,
        preset="script_direction_worker",
        context={"root_trace_id": "root-a"},
    )
    assert value.model == "frozen-direction-model"
    assert value.temperature == 0.2
    assert value.max_iterations == 18


@pytest.mark.asyncio
async def test_role_prompt_resolver_uses_task_pinned_snapshot_and_rechecks_digest() -> None:
    content = "Frozen paragraph worker policy"
    snapshot = ScriptBuildInputSnapshotV1(
        snapshot_id="3",
        script_build_id=9,
        execution_id=1,
        topic_build_id=2,
        topic_id=3,
        topic={},
        account={},
        persona_points=(),
        section_patterns=(),
        strategies=(),
        prompt_manifest=(
            {
                "preset": "script_paragraph_worker",
                "role": "worker",
                "content": content,
                "content_sha256": f"sha256:{sha256(content.encode('utf-8')).hexdigest()}",
            },
        ),
        datasource_manifest={},
        model_manifest={},
        canonical_sha256="sha256:" + "a" * 64,
        created_at=datetime.now(UTC),
    )

    class Snapshots:
        async def get(self, snapshot_id: str, *, script_build_id: int):
            assert (snapshot_id, script_build_id) == ("3", 9)
            return snapshot

    resolver = SnapshotRoleSystemPromptResolver(_Bindings(), Snapshots())
    value = await resolver.resolve(
        role=AgentRole.WORKER,
        preset="script_paragraph_worker",
        context={
            "root_trace_id": "root-a",
            "task_spec": {"context_refs": ["script-build://inputs/3"]},
        },
    )
    assert value is not None
    assert value.content == content

    tampered = replace(
        snapshot,
        prompt_manifest=(
            {
                **snapshot.prompt_manifest[0],
                "content_sha256": "sha256:" + "f" * 64,
            },
        ),
    )

    class TamperedSnapshots:
        async def get(self, _snapshot_id: str, *, script_build_id: int):
            assert script_build_id == 9
            return tampered

    with pytest.raises(ProtocolViolation, match="digest"):
        await SnapshotRoleSystemPromptResolver(_Bindings(), TamperedSnapshots()).resolve(
            role=AgentRole.WORKER,
            preset="script_paragraph_worker",
            context={
                "root_trace_id": "root-a",
                "task_spec": {"context_refs": ["script-build://inputs/3"]},
            },
        )


@pytest.mark.asyncio
async def test_snapshot_tool_hides_on_demand_body_until_explicit_digest_verified_load() -> None:
    strategy_ref = "script-build://strategies/8/versions/2"
    snapshot = ScriptBuildInputSnapshotV1(
        snapshot_id="3",
        script_build_id=9,
        execution_id=1,
        topic_build_id=2,
        topic_id=3,
        topic={},
        account={},
        persona_points=(),
        section_patterns=(),
        strategies=(
            {
                "strategy_id": 7,
                "version": 1,
                "mode": "always_on",
                "content": "always body",
                "content_sha256": "sha256:" + "1" * 64,
            },
            {
                "strategy_id": 8,
                "version": 2,
                "mode": "on_demand",
                "description": "optional",
                "content": "on demand body",
                "content_sha256": canonical_sha256("on demand body").wire,
            },
        ),
        prompt_manifest=(),
        datasource_manifest={},
        model_manifest={},
        canonical_sha256="sha256:" + "a" * 64,
        created_at=datetime.now(UTC),
    )

    class Snapshots:
        async def get(self, _snapshot_id: str, *, script_build_id: int):
            assert script_build_id == 9
            return snapshot

    task = SimpleNamespace(current_spec=SimpleNamespace(context_refs=["script-build://inputs/3"]))
    coordinator = SimpleNamespace(
        task_store=SimpleNamespace(load=lambda _root: _async(SimpleNamespace(tasks={"task": task})))
    )
    gateway = LegacyScriptToolGateway(
        bindings=_Bindings(),
        snapshots=Snapshots(),
        artifacts=SimpleNamespace(),
        coordinator=coordinator,
    )
    context = {"root_trace_id": "root-a", "task_id": "task"}
    summary = await gateway.read_input_snapshot(context)
    assert summary["strategies"][0]["content"] == "always body"
    assert "content" not in summary["strategies"][1]
    loaded = await gateway.load_frozen_strategy(strategy_ref, context)
    assert loaded["content"] == "on demand body"

    tampered = replace(
        snapshot,
        strategies=(
            snapshot.strategies[0],
            {**snapshot.strategies[1], "content": "changed after freeze"},
        ),
    )

    class TamperedSnapshots:
        async def get(self, _snapshot_id: str, *, script_build_id: int):
            assert script_build_id == 9
            return tampered

    gateway.snapshots = TamperedSnapshots()
    with pytest.raises(ProtocolViolation, match="strategy digest"):
        await gateway.load_frozen_strategy(strategy_ref, context)


async def _async(value):
    return value