import yaml
from .schema import AgentConfig
from .llm_client import detect_provider
from pathlib import Path
from typing import Optional, Callable, Any, Dict, List
from .agent import Agent
import os

# Provider → env var mapping
_API_KEY_ENV = {
    "openai": "OPENAI_API_KEY",
    "anthropic": "ANTHROPIC_API_KEY",
    "google": "GOOGLE_API_KEY",
    "ollama": None,  # No key needed
}


def parse_yaml(file_path: Path) -> dict[str, any]:
    """Load YAML and return a nested dict object."""
    with open(file_path, "r", encoding="utf-8") as file:
        data = yaml.safe_load(file)
    return data


def get_agent_config(agent_config_path: Path) -> AgentConfig | None:
    """Load a single agent config YAML and return an AgentConfig."""
    raw = parse_agent_config(agent_config_path)
    if raw is None:
        return None

    # The YAML is either {agent_name: {config}} or flat {config with name}
    agent_name = list(raw.keys())[0] if isinstance(raw, dict) and len(raw) == 1 else raw.get('name')
    config = raw[agent_name] if agent_name in raw else raw

    return AgentConfig.from_yaml(agent_name, config)


def get_agents_configs(agents_dir: Path) -> list[AgentConfig] | None:
    """Load all agent configs from agents.yaml manifest."""
    if not agents_dir.exists():
        raise Exception(f'AGENT CONFIGURATION FILE NOT FOUND: {agents_dir}')

    agents_manifest = parse_agents_config(agents_dir)
    configs: list[AgentConfig] = []

    try:
        for agent in agents_manifest["agents"]:
            agent_config_path = agents_dir / agent["config_path"]
            cfg = get_agent_config(agent_config_path)
            if cfg:
                configs.append(cfg)
    except Exception as e:
        print(f"Error getting agent config: {e}")
        return None

    return configs


def _get_api_key(config: AgentConfig) -> str:
    """Resolve the API key for an agent based on its model provider."""
    provider = config.provider or detect_provider(config.model)
    env_var = _API_KEY_ENV.get(provider)
    if env_var is None:
        return "ollama"  # Local models don't need a key
    key = os.environ.get(env_var)
    if not key:
        raise ValueError(
            f"Agent '{config.name}' uses provider '{provider}' "
            f"but env var {env_var} is not set."
        )
    return key


def agents_configuration(agents_dir: Path) -> dict[str, Agent]:
    """Build all Agent instances from the agents directory."""
    agents = {}
    agents_dir = Path(agents_dir)

    try:
        configs = get_agents_configs(agents_dir)
        if not configs:
            raise Exception('NO AGENTS CONFIGS FOUND')

        for config in configs:
            api_key = _get_api_key(config)
            agent = Agent(agents_dir, config, api_key=api_key)
            print(f'{config.name} AGENT CONFIGURATION SUCCESS ({config.model})')
            agents[config.name] = agent

        return agents
    except Exception as e:
        print(f'AGENTS CONFIGURATION FAILED: {e}')
        raise Exception(f'AGENTS CONFIGURATION FAILED: {e}')


def parse_agent_config(agents_dir: Path):
    print(agents_dir)
    try:
        with open(agents_dir, 'r') as file:
            agent_config = yaml.safe_load(file)
        return agent_config
    except Exception as e:
        print(f'AGENT CONFIGURATION PARSING FAILED: {e}')
        raise Exception(f'AGENT CONFIGURATION PARSING FAILED: {e}')


def parse_agents_config(agents_dir: Path):
    try:
        with open(agents_dir / 'config/agents.yaml', 'r') as file:
            agent_config = yaml.safe_load(file)
        return agent_config
    except Exception as e:
        print(f'AGENT CONFIGURATION PARSING FAILED: {e}')
        raise Exception(f'AGENT CONFIGURATION PARSING FAILED: {e}')