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 (
    CONTEXT_ACCESS_PROTOCOL,
    phase_three_prompt_manifest,
    script_build_prompt_manifest,
)
from script_build_host.domain.errors import ProtocolViolation
from script_build_host.domain.input_snapshot import ScriptBuildInputSnapshotV1
from script_build_host.domain.workbench import protect_model_value
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 []) == {
        "get_current_context",
        "plan_script_tasks",
        "decide_script_task",
        "search_mission_context",
        "read_mission_context",
        "dispatch_script_tasks",
        "validate_attempt",
    }
    compose = get_preset("script_compose_worker")
    assert set(compose.allowed_tools or []) == {"submit_attempt"}
    compare = get_preset("script_candidate_compare_worker")
    assert "save_comparison_candidate" in (compare.allowed_tools or [])
    assert "create_script_paragraph" not in (compare.allowed_tools or [])
    portfolio = get_preset("script_candidate_portfolio_worker")
    assert set(portfolio.allowed_tools or []) == {"submit_attempt"}
    assert "view_frozen_images" not in (
        get_preset("script_retrieval_validator").allowed_tools or []
    )
    assert "view_frozen_images" in (get_preset("script_candidate_validator").allowed_tools or [])
    element_tools = set(get_preset("script_element_set_worker").allowed_tools or [])
    assert "save_script_elements" in element_tools
    assert "create_script_element" not in element_tools
    assert "batch_link_paragraph_elements" not in element_tools

    assert set(get_preset("script_root_worker").allowed_tools or []) == {
        "search_mission_context",
        "read_mission_context",
        "legacy_projection_dry_run",
        "save_root_delivery_manifest",
        "submit_attempt",
    }
    assert set(get_preset("script_root_validator").allowed_tools or []) == {
        "search_mission_context",
        "read_mission_context",
        "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) == 15
    assert len({str(item["preset"]) for item in manifest}) == len(manifest)
    assert {str(item["preset"]) for item in manifest} >= {
        "script_planner",
        "script_context_broker",
        "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)


def test_every_model_role_freezes_one_shared_context_access_protocol() -> None:
    manifest = {
        str(item["preset"]): str(item["content"])
        for item in (*script_build_prompt_manifest(), *phase_three_prompt_manifest())
    }
    deterministic_or_utility = {
        "script_context_broker",
        "script_compose_worker",
        "script_candidate_portfolio_worker",
    }

    for preset, content in manifest.items():
        if preset in deterministic_or_utility:
            assert not content.startswith(CONTEXT_ACCESS_PROTOCOL)
        else:
            assert content.startswith(CONTEXT_ACCESS_PROTOCOL), preset
            assert content.count("## Context Access Protocol") == 1

    assert "active_has_more" in manifest["script_planner"]
    for preset in (
        "script_pattern_retrieval_worker",
        "script_decode_retrieval_worker",
        "script_external_retrieval_worker",
        "script_knowledge_retrieval_worker",
    ):
        assert "read_mission_context` 的 cursor" in manifest[preset]
        assert "不计作第二次" in manifest[preset]
    assert "Direction、CandidatePortfolio" in manifest["script_root_worker"]
    assert "required_handles" in manifest["script_candidate_validator"]
    assert "required_handles" in manifest["script_root_validator"]


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()
    }
    descriptions = {
        item["function"]["name"]: item["function"]["description"] 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["search_mission_context"]["required"] == ["collection"]
    assert schemas["read_mission_context"]["required"] == ["handle"]
    assert "sufficiency" in descriptions["search_mission_context"]
    assert "exhausted=true" in descriptions["read_mission_context"]
    assert "load_images" not in schemas
    assert "load_frozen_strategy" not in schemas
    assert "read_workbench_detail" not in schemas
    assert schemas["save_script_elements"]["required"] == [
        "elements",
        "links",
        "expected_state_revision",
    ]
    assert schemas["save_script_elements"]["properties"]["links"]["items"]["required"] == [
        "paragraph_target_key",
        "element_client_keys",
    ]
    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",
    }


def test_agent_visible_write_schemas_never_expose_mechanical_identity_fields() -> 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()
    }
    write_tools = {
        "plan_script_tasks",
        "decide_script_task",
        "save_direction_candidate",
        "save_script_paragraphs",
        "save_script_elements",
        "save_comparison_candidate",
        "save_root_delivery_manifest",
        "submit_validation",
    }
    forbidden = {
        "paragraph_id",
        "element_id",
        "branch_id",
        "artifact_ref",
        "evidence_refs",
        "write_scope",
        "output_schema",
        "compose_order",
        "context_refs",
    }

    def property_names(value: object) -> set[str]:
        if isinstance(value, dict):
            names = (
                set(value.get("properties", {}))
                if isinstance(value.get("properties"), dict)
                else set()
            )
            return names | set().union(*(property_names(item) for item in value.values()))
        if isinstance(value, list):
            return set().union(*(property_names(item) for item in value))
        return set()

    for name in write_tools:
        assert forbidden.isdisjoint(property_names(schemas[name])), name

    register_script_presets()
    forbidden_tools = {
        "read_input_snapshot",
        "read_attempt_workspace",
        "read_active_frontier",
        "read_pinned_candidates",
        "read_accepted_artifact",
        "read_workbench_detail",
        "load_frozen_strategy",
        "load_images",
        "create_script_paragraph",
        "create_script_element",
        "update_script_element",
        "batch_link_paragraph_elements",
        "append_paragraph_atoms",
        "delete_paragraph_atom",
    }
    for preset_name in (
        "script_direction_worker",
        "script_structure_worker",
        "script_paragraph_worker",
        "script_element_set_worker",
        "script_candidate_compare_worker",
        "script_compose_worker",
        "script_candidate_portfolio_worker",
        "script_root_worker",
    ):
        assert forbidden_tools.isdisjoint(get_preset(preset_name).allowed_tools or []), preset_name


def test_model_visible_values_replace_storage_identity_and_uri_fields() -> None:
    protected = protect_model_value(
        {
            "uri": "script-build://artifact-versions/9",
            "url": "https://private.example/source",
            "raw_artifact_ref": "script-build://raw-artifacts/sha256/secret",
            "artifact_version_id": 9,
            "branch_id": 7,
            "strategy_id": 11,
            "topic_id": 12,
            "topic_build_id": 13,
            "execution_id": 14,
            "content_sha256": "sha256:secret",
            "goal_id": "goal-allowed",
            "decision_id": "decision-allowed",
        },
        namespace="surface-test",
    )
    encoded = str(protected)

    assert "script-build://" not in encoded
    assert "https://" not in encoded
    assert "sha256:" not in encoded
    assert "artifact_version_id" not in protected
    assert "branch_id" not in protected
    assert "strategy_id" not in protected
    assert "topic_id" not in protected
    assert "topic_build_id" not in protected
    assert "execution_id" not in protected
    assert protected["goal_id"] == "goal-allowed"
    assert protected["decision_id"] == "decision-allowed"


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"]},
            },
        )


async def _async(value):
    return value