from __future__ import annotations

from datetime import UTC, datetime
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.domain.input_snapshot import ScriptBuildInputSnapshotV1
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",
    }
    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 []))


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"]


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