"""Task-driven tool assembler.

这层和 SkillAssembler 的位置类似：它不执行工具，只决定本轮 run 应该把哪些
tool schema 暴露给模型。
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING

from beaver.engine.context import SkillContext
from beaver.foundation.embedding import EmbeddingRetriever
from beaver.tools.base import ToolSpec
from beaver.tools.registry import ToolRegistry

if TYPE_CHECKING:
    from beaver.engine.providers.runtime import ProviderRuntime
    from beaver.skills.catalog.loader import SkillsLoader


class ToolAssembler:
    """Use skill hints and embedding retrieval to select run-scoped tools."""

    def __init__(
        self,
        *,
        retriever: EmbeddingRetriever | None = None,
        always_tool_names: Sequence[str] | None = None,
    ) -> None:
        self.retriever = retriever or EmbeddingRetriever()
        self.always_tool_names = tuple(always_tool_names or ("memory", "session_search", "skill_view"))

    async def assemble(
        self,
        *,
        task_description: str,
        registry: ToolRegistry,
        skills_loader: SkillsLoader | None = None,
        activated_skills: Sequence[SkillContext] | None = None,
        embedding_runtime: ProviderRuntime | None = None,
        top_k: int = 10,
    ) -> list[ToolSpec]:
        """Return selected tool specs for the current run.

        Selection order is intentionally deterministic:
        1. always tools from config/spec
        2. tools explicitly declared by activated skills
        3. embedding top-k tools for the task
        """

        selected: list[ToolSpec] = []
        selected_names: set[str] = set()

        def add_specs(specs: Sequence[ToolSpec]) -> None:
            for spec in specs:
                if spec.name in selected_names:
                    continue
                selected.append(spec)
                selected_names.add(spec.name)

        add_specs(registry.list_always_specs())
        add_specs(registry.get_specs(self.always_tool_names))

        skill_tool_names = self._collect_skill_tool_names(
            skills_loader=skills_loader,
            activated_skills=activated_skills or (),
        )
        add_specs(registry.get_specs(skill_tool_names))

        candidates = [
            spec.to_embedding_candidate()
            for spec in registry.list_specs()
            if spec.name not in selected_names
        ]
        retrieved = await self.retriever.retrieve(
            query=task_description,
            candidates=candidates,
            top_k=top_k,
            api_key=embedding_runtime.api_key if embedding_runtime is not None else None,
            api_base=embedding_runtime.api_base if embedding_runtime is not None else None,
            model=embedding_runtime.model if embedding_runtime is not None else None,
            extra_headers=embedding_runtime.extra_headers if embedding_runtime is not None else None,
            timeout_seconds=(
                embedding_runtime.request_timeout_seconds if embedding_runtime is not None else None
            ),
            fallback_top_k=top_k,
        )
        add_specs(registry.get_specs([item["name"] for item in retrieved]))
        return selected

    @staticmethod
    def _collect_skill_tool_names(
        *,
        skills_loader: SkillsLoader | None,
        activated_skills: Sequence[SkillContext],
    ) -> list[str]:
        if skills_loader is None or not activated_skills:
            return []

        result: list[str] = []
        for skill in activated_skills:
            names = list(skill.tool_hints) if getattr(skill, "tool_hints", None) else skills_loader.get_skill_tool_hints(skill.name)
            for name in names:
                if name not in result:
                    result.append(name)
        return result