"""Plugin system — head registry for custom heads.

Provides a registry-based architecture for dynamically registering,
discovering, and creating head agents. Replaces the hardcoded head
creation in ``agents/heads/__init__.py`` with an extensible system.

Usage:

    from fusionagi.agents.head_registry import HeadRegistry

    registry = HeadRegistry()

    # Built-in heads are pre-registered
    head = registry.create("logic")

    # Register a custom head
    @registry.register_factory("my_domain")
    def create_my_head(adapter, **kwargs):
        return HeadAgent(head_id=HeadId.LOGIC, role="My Domain", ...)

    # Discover all available heads
    registry.list_heads()
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable

from fusionagi._logger import logger
from fusionagi.adapters.base import LLMAdapter
from fusionagi.agents.head_agent import HeadAgent
from fusionagi.prompts.heads import get_head_prompt
from fusionagi.reasoning.native import NativeReasoningProvider
from fusionagi.schemas.head import HeadId


@dataclass
class HeadSpec:
    """Specification for a registered head type."""

    head_id: str
    role: str
    objective: str
    factory: Callable[..., HeadAgent]
    description: str = ""
    tags: list[str] = field(default_factory=list)
    builtin: bool = True


class HeadRegistry:
    """Extensible registry for head agent types.

    Pre-registers all 11 built-in Dvādaśa content heads on creation.
    Custom heads can be added via ``register()`` or ``register_factory()``.
    """

    def __init__(self, *, auto_register_builtins: bool = True) -> None:
        self._specs: dict[str, HeadSpec] = {}
        if auto_register_builtins:
            self._register_builtins()

    def _register_builtins(self) -> None:
        """Register all built-in Dvādaśa content heads."""
        role_map: dict[HeadId, tuple[str, str]] = {
            HeadId.LOGIC: ("Logic", "Correctness, contradictions, formal checks"),
            HeadId.RESEARCH: ("Research", "Retrieval, source quality, citations"),
            HeadId.SYSTEMS: ("Systems", "Architecture, dependencies, scalability"),
            HeadId.STRATEGY: ("Strategy", "Roadmap, prioritization, tradeoffs"),
            HeadId.PRODUCT: ("Product/UX", "Interaction design, user flows"),
            HeadId.SECURITY: ("Security", "Threats, auth, secrets, abuse vectors"),
            HeadId.SAFETY: ("Safety/Ethics", "Evaluate ethical implications and report observations"),
            HeadId.RELIABILITY: ("Reliability", "SLOs, failover, load testing, observability"),
            HeadId.COST: ("Cost/Performance", "Token budgets, caching, model routing"),
            HeadId.DATA: ("Data/Memory", "Schemas, privacy, retention, personalization"),
            HeadId.DEVEX: ("DevEx", "CI/CD, testing strategy, local tooling"),
        }

        for head_id, (role, objective) in role_map.items():
            self._register_builtin_head(head_id, role, objective)

    def _register_builtin_head(
        self, head_id: HeadId, role: str, objective: str
    ) -> None:
        """Register a single built-in head."""

        def factory(
            adapter: LLMAdapter | None = None,
            tool_permissions: list[str] | None = None,
            reasoning_provider: NativeReasoningProvider | None = None,
            use_native_reasoning: bool = True,
            _hid: HeadId = head_id,
            _role: str = role,
            _obj: str = objective,
            **kwargs: Any,
        ) -> HeadAgent:
            provider = reasoning_provider
            if provider is None and use_native_reasoning and adapter is None:
                provider = NativeReasoningProvider()

            return HeadAgent(
                head_id=_hid,
                role=_role,
                objective=_obj,
                system_prompt=get_head_prompt(_hid),
                adapter=adapter,
                tool_permissions=tool_permissions,
                reasoning_provider=provider,
            )

        self._specs[head_id.value] = HeadSpec(
            head_id=head_id.value,
            role=role,
            objective=objective,
            factory=factory,
            description=f"Built-in {role} head",
            tags=["builtin", "dvadasa"],
            builtin=True,
        )

    def register(
        self,
        head_id: str,
        role: str,
        objective: str,
        factory: Callable[..., HeadAgent],
        *,
        description: str = "",
        tags: list[str] | None = None,
    ) -> None:
        """Register a custom head type.

        Args:
            head_id: Unique identifier for the head.
            role: Head's role name.
            objective: What the head does.
            factory: Callable that creates a HeadAgent.
            description: Human-readable description.
            tags: Optional tags for discovery.
        """
        if head_id in self._specs:
            logger.warning(
                "Overwriting existing head registration",
                extra={"head_id": head_id},
            )

        self._specs[head_id] = HeadSpec(
            head_id=head_id,
            role=role,
            objective=objective,
            factory=factory,
            description=description,
            tags=tags or [],
            builtin=False,
        )
        logger.info("Custom head registered", extra={"head_id": head_id, "role": role})

    def register_factory(
        self,
        head_id: str,
        *,
        role: str = "",
        objective: str = "",
        description: str = "",
        tags: list[str] | None = None,
    ) -> Callable[[Callable[..., HeadAgent]], Callable[..., HeadAgent]]:
        """Decorator to register a head factory function.

        Args:
            head_id: Unique identifier.
            role: Head's role name.
            objective: What the head does.
            description: Human-readable description.
            tags: Optional tags.

        Returns:
            Decorator function.
        """

        def decorator(fn: Callable[..., HeadAgent]) -> Callable[..., HeadAgent]:
            self.register(
                head_id=head_id,
                role=role or head_id.replace("_", " ").title(),
                objective=objective or fn.__doc__ or "",
                factory=fn,
                description=description,
                tags=tags,
            )
            return fn

        return decorator

    def create(
        self,
        head_id: str,
        adapter: LLMAdapter | None = None,
        **kwargs: Any,
    ) -> HeadAgent:
        """Create a head agent by ID.

        Args:
            head_id: Registered head identifier.
            adapter: Optional LLM adapter.
            **kwargs: Additional arguments passed to factory.

        Returns:
            Created HeadAgent.

        Raises:
            KeyError: If head_id is not registered.
        """
        if head_id not in self._specs:
            raise KeyError(
                f"Head '{head_id}' not registered. "
                f"Available: {', '.join(sorted(self._specs.keys()))}"
            )
        spec = self._specs[head_id]
        return spec.factory(adapter=adapter, **kwargs)

    def create_all(
        self,
        adapter: LLMAdapter | None = None,
        *,
        include_tags: list[str] | None = None,
        exclude_tags: list[str] | None = None,
        **kwargs: Any,
    ) -> dict[str, HeadAgent]:
        """Create all registered heads (optionally filtered by tags).

        Args:
            adapter: Optional LLM adapter.
            include_tags: Only create heads matching these tags.
            exclude_tags: Skip heads matching these tags.
            **kwargs: Additional arguments.

        Returns:
            Dict of head_id -> HeadAgent.
        """
        heads: dict[str, HeadAgent] = {}
        for hid, spec in self._specs.items():
            if include_tags and not any(t in spec.tags for t in include_tags):
                continue
            if exclude_tags and any(t in spec.tags for t in exclude_tags):
                continue
            heads[hid] = spec.factory(adapter=adapter, **kwargs)
        return heads

    def list_heads(self) -> list[dict[str, Any]]:
        """List all registered heads.

        Returns:
            List of head specifications.
        """
        return [
            {
                "head_id": spec.head_id,
                "role": spec.role,
                "objective": spec.objective,
                "description": spec.description,
                "tags": spec.tags,
                "builtin": spec.builtin,
            }
            for spec in self._specs.values()
        ]

    def get_spec(self, head_id: str) -> HeadSpec | None:
        """Get the spec for a registered head."""
        return self._specs.get(head_id)

    def unregister(self, head_id: str) -> bool:
        """Remove a head registration.

        Args:
            head_id: Head to remove.

        Returns:
            True if removed, False if not found.
        """
        if head_id in self._specs:
            del self._specs[head_id]
            return True
        return False

    def broadcast_ethical_feedback(
        self,
        heads: dict[str, Any],
        feedback: dict[str, Any],
    ) -> None:
        """Broadcast ethical feedback to all active heads.

        Args:
            heads: Dict of head_id -> HeadAgent instances.
            feedback: Ethical feedback data.
        """
        for hid, head in heads.items():
            if hasattr(head, "on_ethical_feedback"):
                head.on_ethical_feedback(feedback)

    def broadcast_consequence(
        self,
        heads: dict[str, Any],
        consequence: dict[str, Any],
    ) -> None:
        """Broadcast consequence data to all active heads.

        Args:
            heads: Dict of head_id -> HeadAgent instances.
            consequence: Consequence data.
        """
        for hid, head in heads.items():
            if hasattr(head, "on_consequence"):
                head.on_consequence(consequence)

    @property
    def registered_count(self) -> int:
        """Number of registered heads."""
        return len(self._specs)


# Global default registry
_default_registry: HeadRegistry | None = None


def get_default_registry() -> HeadRegistry:
    """Get or create the default global head registry."""
    global _default_registry  # noqa: PLW0603
    if _default_registry is None:
        _default_registry = HeadRegistry()
    return _default_registry


__all__ = [
    "HeadRegistry",
    "HeadSpec",
    "get_default_registry",
]