"""Utilities for loading agents from files."""

from __future__ import annotations

import importlib.util
import logging
from pathlib import Path

from yera.agents.decorator import AgentFunctionWrapper
from yera.utils.path_utils import path_to_module_name


def load_agent(
    file_path: Path | str, agent_identifier: str | None = None
) -> AgentFunctionWrapper:
    """Load a Python file and extract the AgentFunctionWrapper instance with the given identifier.

    Args:
        file_path: Path to the Python file containing one or more agents.
            Can be a Path object or string.
        agent_identifier: Identifier of the agent to load. If None, the first agent is loaded.

    Returns:
        AgentFunctionWrapper instance.

    Raises:
        ValueError: If no agents are found in the file.
        ValueError: If multiple agents are found in the file and no identifier is provided.
        ValueError: If the specified agent identifier is not found in the file.
        ValueError: If the specified agent identifier does not match the identifier of the loaded agent.
    """
    agents = load_agents(file_path)

    if not agents:
        raise ValueError(f"No agents found in {file_path}")

    if agent_identifier is None:
        if len(agents) > 1:
            raise ValueError(
                f"Multiple agents found in {file_path}, please specify an agent identifier"
            )
        return agents[0]

    matching_agents = [
        agent for agent in agents if agent.metadata.identifier == agent_identifier
    ]
    if not matching_agents:
        raise ValueError(f"Agent with ID {agent_identifier} not found in {file_path}")
    if len(matching_agents) > 1:
        raise ValueError(
            f"Multiple agents with ID {agent_identifier} found in {file_path}"
        )
    return matching_agents[0]


def load_agents(file_path: Path | str) -> list[AgentFunctionWrapper]:
    """Load a Python file and extract all AgentFunctionWrapper instances.

    This helper looks for objects in the module that are instances of
    AgentFunctionWrapper (created by the @yr.agent decorator).

    Args:
        file_path: Path to the Python file containing one or more agents.
            Can be a Path object or string.

    Returns:
        Dictionary mapping agent identifiers to AgentFunctionWrapper instances.

    Raises:
        FileNotFoundError: If the file does not exist.
    """
    path = Path(file_path) if isinstance(file_path, str) else file_path
    if not path.exists():
        raise FileNotFoundError(f"Python file not found: {path}")

    # Load the module from the given file path
    module_name = path_to_module_name(path)
    spec = importlib.util.spec_from_file_location(module_name, path)
    if spec is None or spec.loader is None:
        raise FileNotFoundError(f"Could not load Python module from: {path}")

    module = importlib.util.module_from_spec(spec)
    # Exec here is fine because we are only allowing users to run via the CLI locally,
    # so there is no risk of arbitrary remote code execution.
    # Deployment will require tracing, which will NOT exec the python directly.
    spec.loader.exec_module(module)

    # Collect all AgentFunctionWrapper instances defined in the module
    return [
        value
        for value in vars(module).values()
        if isinstance(value, AgentFunctionWrapper)
    ]


def discover_agents(directory: Path | str) -> list[AgentFunctionWrapper]:
    """Discover all agents in a directory.

    Scans the directory for .py files and loads all agents from each file.

    This function underpins the `/api/agents` endpoint exposed by the
    stream server: the React UI queries `/api/agents` to build the home
    page and to resolve `agent_id` strings used in `/agents/{id}` routes.
    Running ``yera dev`` in different directories will therefore expose
    different sets of agents to the *same* static UI bundle.

    Args:
        directory: Directory to scan for agent files.

    Returns:
        List of AgentFunctionWrapper instances.

    Raises:
        FileNotFoundError: If the directory does not exist.
    """
    dir_path = Path(directory) if isinstance(directory, str) else directory
    if not dir_path.exists():
        raise FileNotFoundError(f"Directory not found: {dir_path}")
    if not dir_path.is_dir():
        raise ValueError(f"Path is not a directory: {dir_path}")

    all_agents = []
    for py_file in dir_path.rglob("*.py"):
        try:
            all_agents.extend(load_agents(py_file))
        # Ignore PERF203 because I/O file loading is the real bottleneck.
        except FileNotFoundError:  # noqa: PERF203
            logging.warning(f"Failed to load agents from {py_file}")
            continue
    return all_agents
