"""Policy guardrails for swarms-generated agent team plans."""

from __future__ import annotations

from typing import Any


class SwarmsPolicy:
    """Clamp AutoSwarmBuilder output before nanobot executes it."""

    allowed_swarm_types = {
        # Keep this list to swarms that consume the provided nanobot agent adapters.
        "GroupChat",
        "SequentialWorkflow",
        "ConcurrentWorkflow",
        "AgentRearrange",
        "MixtureOfAgents",
        "HierarchicalSwarm",
    }

    def __init__(self, *, max_agents: int = 4, max_loops: int = 3) -> None:
        self.max_agents = max(1, max_agents)
        self.max_loops = max(1, max_loops)

    def validate_auto_config(self, raw_config: dict[str, Any]) -> dict[str, Any]:
        config = self._plain_dict(raw_config)

        swarm_type = str(
            config.get("swarm_type")
            or config.get("type")
            or config.get("architecture")
            or "GroupChat"
        )
        if swarm_type not in self.allowed_swarm_types:
            swarm_type = "GroupChat"
        config["swarm_type"] = swarm_type

        agents = list(config.get("agents") or [])[: self.max_agents]
        config["agents"] = [self._sanitize_agent_spec(item) for item in agents]
        config["max_loops"] = min(max(1, int(config.get("max_loops") or 2)), self.max_loops)

        # AutoSwarmBuilder may suggest structure, not grant capabilities.
        config.pop("tools", None)
        config.pop("mcp_url", None)
        config.pop("mcp_urls", None)
        config.pop("llm_api_key", None)
        config.pop("api_key", None)
        return config

    def _plain_dict(self, raw_config: Any) -> dict[str, Any]:
        if isinstance(raw_config, dict):
            return dict(raw_config)
        model_dump = getattr(raw_config, "model_dump", None)
        if callable(model_dump):
            payload = model_dump()
            return dict(payload) if isinstance(payload, dict) else {}
        dict_method = getattr(raw_config, "dict", None)
        if callable(dict_method):
            payload = dict_method()
            return dict(payload) if isinstance(payload, dict) else {}
        return {}

    def _sanitize_agent_spec(self, item: Any) -> dict[str, Any]:
        spec = self._plain_dict(item)
        return {
            "agent_name": str(spec.get("agent_name") or spec.get("name") or "specialist"),
            "description": str(spec.get("description") or spec.get("agent_description") or ""),
            "system_prompt": str(spec.get("system_prompt") or "")[:4000],
            "role": str(spec.get("role") or "worker"),
        }