"""Agent Registry — loads agent configurations from YAML and instantiates agents."""

from __future__ import annotations

import logging
import re
from pathlib import Path
from typing import TYPE_CHECKING

import yaml
from pydantic import ValidationError

from vastian.agents.base import AgentConfig, BaseAgent
from vastian.bus.bus import MessageBus

if TYPE_CHECKING:
    from vastian.knowledge.store import KnowledgeStore
    from vastian.providers.llm_provider import BaseLLMProvider

logger = logging.getLogger(__name__)


class AgentRegistry:
    """
    Central registry that:
    1. Loads agent definitions from YAML files
    2. Validates configurations
    3. Instantiates agent runtime objects
    4. Provides lookup by ID, role, team, or reporting chain
    """

    def __init__(self) -> None:
        self._configs: dict[str, AgentConfig] = {}
        self._agents: dict[str, BaseAgent] = {}
        self._template_vars: dict[str, str] = {}

    # ── Loading ──────────────────────────────────────

    def load_from_directory(
        self,
        agents_dir: Path,
        template_vars: dict[str, str] | None = None,
    ) -> int:
        """Load all YAML agent definitions from a directory. Returns count loaded.

        Template variables (e.g. ``{"user_name": "Rowan"}``) are substituted into
        ``{{var}}`` placeholders in the YAML text before parsing.
        """
        self._template_vars = template_vars or {}
        if not agents_dir.exists():
            logger.warning("Agents directory not found: %s", agents_dir)
            return 0

        count = 0
        for yaml_file in sorted(agents_dir.glob("*.yaml")):
            try:
                self._load_yaml(yaml_file)
                count += 1
            except (yaml.YAMLError, ValidationError, TypeError) as e:
                logger.error("Failed to load agent config %s: %s", yaml_file.name, e)

        logger.info("Loaded %d agent configurations from %s", count, agents_dir)
        return count

    def _load_yaml(self, path: Path) -> AgentConfig:
        """Load and validate a single agent YAML file.

        Supports ``{{var_name}}`` template placeholders which are replaced from
        ``self._template_vars`` before YAML parsing.
        """
        text = path.read_text(encoding="utf-8")
        # Substitute {{var}} placeholders
        if self._template_vars:
            for key, value in self._template_vars.items():
                text = re.sub(r"\{\{\s*" + re.escape(key) + r"\s*\}\}", value, text)
        # Strip any remaining {{...}} placeholders so YAML doesn't choke on braces
        text = re.sub(r"\{\{[^}]*\}\}", '""', text)
        raw = yaml.safe_load(text)
        if raw is None:
            raise yaml.YAMLError(f"Empty or invalid YAML in {path.name}")
        config = AgentConfig(**raw)
        self._configs[config.agent_id] = config
        logger.debug("Loaded config: %s", config.display_name)
        return config

    # ── Instantiation ────────────────────────────────

    def instantiate_all(
        self,
        bus: MessageBus,
        knowledge_store: KnowledgeStore | None = None,
        llm_provider: BaseLLMProvider | None = None,
    ) -> dict[str, BaseAgent]:
        """Create agent instances for all loaded configurations."""
        for agent_id, config in self._configs.items():
            agent = BaseAgent(
                config=config,
                bus=bus,
                knowledge_store=knowledge_store,
                llm_provider=llm_provider,
            )
            self._agents[agent_id] = agent

        # Register teams on the bus
        self._register_teams(bus)

        logger.info("Instantiated %d agents", len(self._agents))
        return dict(self._agents)

    def _register_teams(self, bus: MessageBus) -> None:
        """Build team memberships from agent configs and register on bus."""
        teams: dict[str, list[str]] = {}
        for config in self._configs.values():
            for team_id in config.teams:
                teams.setdefault(team_id, []).append(config.agent_id)
        for team_id, members in teams.items():
            bus.register_team(team_id, members)

    # ── Lookup ───────────────────────────────────────

    def get_config(self, agent_id: str) -> AgentConfig | None:
        return self._configs.get(agent_id)

    def get_agent(self, agent_id: str) -> BaseAgent | None:
        return self._agents.get(agent_id)

    def get_all_agents(self) -> dict[str, BaseAgent]:
        return dict(self._agents)

    def get_direct_reports(self, agent_id: str) -> list[str]:
        """Get agent IDs that report to the given agent."""
        config = self._configs.get(agent_id)
        return config.direct_reports if config else []

    def get_reporting_chain(self, agent_id: str) -> list[str]:
        """Get the full chain of command from agent up to the top."""
        chain: list[str] = []
        current = agent_id
        visited: set[str] = set()
        while current and current not in visited:
            visited.add(current)
            config = self._configs.get(current)
            if config and config.reports_to:
                chain.append(config.reports_to)
                current = config.reports_to
            else:
                break
        return chain

    def get_team_members(self, team_id: str) -> list[str]:
        """Get all agent IDs belonging to a team."""
        return [config.agent_id for config in self._configs.values() if team_id in config.teams]

    @property
    def agent_ids(self) -> list[str]:
        return list(self._configs.keys())

    @property
    def all_configs(self) -> list[AgentConfig]:
        return list(self._configs.values())
