from __future__ import annotations

import json
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.domain.workbench import strategy_handle
from script_build_host.tools.gateway import LegacyScriptToolGateway, _project_snapshot
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",
        "read_workbench_detail",
        "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 []) == {
        "read_workbench_detail",
        "legacy_projection_dry_run",
        "save_root_delivery_manifest",
        "submit_attempt",
    }
    assert set(get_preset("script_root_validator").allowed_tools or []) == {
        "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_handle"]
    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",
        "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_input_snapshot_projection_is_bounded_and_keeps_business_inputs() -> None:
    persona = [
        {
            "列": ("实质", "形式", "作用", "感受")[index % 4],
            "点名称": f"point-{index}",
            "人设权重分": str(1 - index / 100),
            "帖子覆盖率": "0.8",
            "分类路径": ["/account/style"],
        }
        for index in range(87)
    ]
    payload = {
        "snapshot_id": "50",
        "script_build_id": 900049,
        "canonical_sha256": "sha256:" + "a" * 64,
        "account": {"account_name": "account"},
        "topic": {
            "topic": {"result": "gym story", "topic_direction": "satirical comic"},
            "composition_items": [
                {
                    "id": index,
                    "point_type": "关键点",
                    "dimension": "形式",
                    "element_name": f"item-{index}",
                    "note": "useful",
                    "created_at": "unused metadata",
                    "is_active": True,
                }
                for index in range(30)
            ],
        },
        "persona_points": persona,
        "section_patterns": [{"模式名称": "parallel atlas"}],
        "strategies": [],
        "prompt_manifest": [{"content": "x" * 100_000}],
        "model_manifest": {"secret": "x" * 100_000},
        "datasource_manifest": {"endpoint": "x" * 100_000},
    }
    projected = _project_snapshot(payload, view="element-set")
    encoded = json.dumps(projected, ensure_ascii=False)

    assert projected["view"] == "element-set"
    assert projected["persona_points_total"] == 87
    assert len(projected["persona_points"]) == 12
    assert projected["section_patterns"] == [{"模式名称": "parallel atlas"}]
    assert "prompt_manifest" not in projected
    assert "model_manifest" not in projected
    assert "datasource_manifest" not in projected
    assert "created_at" not in projected["topic"]["composition_items"][0]
    assert len(encoded) < 20_000


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:
    on_demand_body = "strategy detail " * 1_200
    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]
    selected_handle = strategy_handle(snapshot.snapshot_id, 1, snapshot.strategies[1])
    loaded = await gateway.load_frozen_strategy(selected_handle, context)
    assert loaded["content_fragment"] == on_demand_body[:7_000]
    assert loaded["next_cursor"] == 7_000
    second = await gateway.load_frozen_strategy(
        selected_handle, context, cursor=loaded["next_cursor"]
    )
    assert second["content_fragment"] == on_demand_body[7_000:14_000]
    assert "content_sha256" not in loaded["metadata"]

    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(selected_handle, context)


async def _async(value):
    return value