"""Agent Registry and factory for Fluxibly.

Mirrors the LLM Registry pattern. Maps agent class names → agent classes,
and provides create_agent() to instantiate the correct class from config.
"""

from __future__ import annotations

import importlib
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from fluxibly.agent.base import AgentConfig, BaseAgent

_agent_registry: dict[str, type] = {}


def register_agent(name: str, cls: type) -> None:
    """Register an agent class.

    Built-in agents (e.g. SimpleAgent) are auto-registered at import time.
    Call this only for custom agent classes.

    Args:
        name: Class name (e.g. ``"SimpleAgent"``, ``"MyCustomAgent"``).
        cls: BaseAgent subclass.
    """
    _agent_registry[name] = cls


def create_agent(agent: str | AgentConfig) -> BaseAgent:
    """Factory function to create an agent from a config path/name or AgentConfig.

    Resolution:
        1. If ``agent`` is a ``str`` → load config via ``load_agent_config()``
        2. Read ``config.agent_class`` to determine which class to instantiate
        3. Resolve the class name via the registry or dynamic import
        4. Instantiate and return

    ``agent_class`` resolution order:
        1. Registered name (e.g. ``"SimpleAgent"``) → use registered class
        2. ``"module.path::ClassName"`` → import module, get class
        3. Not found → raise ``ValueError``

    Args:
        agent: Config path/name (str) or AgentConfig object.

    Returns:
        Instantiated BaseAgent subclass.

    Raises:
        FileNotFoundError: If config path not found.
        ValueError: If agent_class cannot be resolved.
    """
    from fluxibly.agent.base import AgentConfig
    from fluxibly.config.loader import load_agent_config

    # Step 1: Get the AgentConfig
    if isinstance(agent, str):
        config = load_agent_config(agent)
    elif isinstance(agent, AgentConfig):
        config = agent
    else:
        raise TypeError(
            f"create_agent() expects str or AgentConfig, got {type(agent)}"
        )

    # Step 2: Resolve agent_class to a Python class
    agent_class_ref = config.agent_class
    cls = _resolve_agent_class(agent_class_ref)

    # Step 3: Instantiate
    return cls(config)


def _resolve_agent_class(class_ref: str) -> type:
    """Resolve an agent_class string to a Python class.

    Resolution order:
        1. Registered name → direct lookup
        2. ``"module.path::ClassName"`` → dynamic import
        3. Not found → raise ValueError
    """
    # 1. Check registry
    if class_ref in _agent_registry:
        return _agent_registry[class_ref]

    # 2. Dynamic import: "module.path::ClassName"
    if "::" in class_ref:
        module_path, class_name = class_ref.rsplit("::", 1)
        try:
            module = importlib.import_module(module_path)
        except ImportError as e:
            raise ValueError(
                f"Could not import agent module '{module_path}': {e}"
            ) from e

        cls = getattr(module, class_name, None)
        if cls is None:
            raise ValueError(
                f"Module '{module_path}' does not have class '{class_name}'"
            )
        return cls

    # 3. Not found
    available = list(_agent_registry.keys())
    raise ValueError(
        f"Unknown agent class: '{class_ref}'. "
        f"Registered: {available}. "
        f"For custom classes, use 'module.path::ClassName' format."
    )


def list_agents() -> list[str]:
    """List all registered agent class names."""
    return list(_agent_registry.keys())


def _auto_register() -> None:
    """Auto-register built-in agent implementations."""
    from fluxibly.agent.simple_agent import SimpleAgent

    register_agent("SimpleAgent", SimpleAgent)

    try:
        from fluxibly.agent.agent_template import AgentTemplate

        register_agent("AgentTemplate", AgentTemplate)
    except ImportError:
        pass

    try:
        from fluxibly.agent.multimodal_agent import MultimodalAgent

        register_agent("MultimodalAgent", MultimodalAgent)
    except ImportError:
        pass


_auto_register()
