import importlib.util
from typing import Literal

import yaml
from pydantic import ValidationError

from bench_af._abstract.configs import RunDetectorConfig, ValidateModelConfig
from bench_af._abstract.objects import Detector, Environment, ModelOrganism

from .list import _list_objects
from .root_location import get_bench_af_root

def load_model_organism(name: str) -> ModelOrganism:
    """Load the model organism with the specified file name."""
    model_organism = _load_object(name, "models")
    assert isinstance(model_organism, ModelOrganism), (
        f"Model organism {name} is not a ModelOrganism"
    )
    return model_organism


def load_detector(name: str) -> Detector:
    """Load the detector with the specified file name."""
    detector = _load_object(name, "detectors")
    assert isinstance(detector, Detector), f"Detector {name} is not a Detector"
    return detector


def load_environment(name: str) -> Environment:
    """Load an environment by name."""
    environment = _load_object(name, "environments")
    assert isinstance(environment, Environment), (
        f"Environment {name} is not an Environment"
    )
    return environment


def _load_object(
    name: str, category: Literal["models", "detectors", "environments"]
) -> ModelOrganism | Detector | Environment:
    """Load the object with the specified name and category."""
    if name not in _list_objects(category):
        raise ValueError(
            f"{category.capitalize()} {name} is not in the registry")

    root_path = get_bench_af_root()
    module_path = root_path / category / name / "main.py"
    module_name = f"{category}.{name}.main"
    spec = importlib.util.spec_from_file_location(module_name, module_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    
    assert hasattr(module, "get"), (
        f"{category.capitalize()} {name} must have a `get` attribute"
    )
    get = module.get
    assert callable(get), (
        f"{category.capitalize()} {name} must have a `get` attribute that is callable"
    )
    obj = get()
    assert isinstance(obj, ModelOrganism | Detector | Environment), (
        f"{category.capitalize()} {name} must return a ModelOrganism, Detector, or Environment"
    )
    return obj


def load_config(
    path: str, config_type: Literal["validate-model", "run-detector"]
) -> ValidateModelConfig | RunDetectorConfig:
    """Load and validate the config from the specified file using Pydantic models.

    Args:
        path: Path to the YAML config file
        config_type: Type of config to validate

    Returns:
        Validated configuration as a dictionary

    Raises:
        ValueError: If config is invalid or file cannot be read
        ValidationError: If config doesn't match the expected schema
    """
    with open(path) as f:
        raw_config = yaml.safe_load(f)

    # Validate config using appropriate Pydantic model
    try:
        validated_config: ValidateModelConfig | RunDetectorConfig
        match config_type:
            case "validate-model":
                validated_config = ValidateModelConfig(**raw_config)
            case "run-detector":
                validated_config = RunDetectorConfig(**raw_config)

        return validated_config

    except ValidationError as e:
        # Format validation errors nicely
        error_details = []
        for error in e.errors():
            field = " -> ".join(str(loc) for loc in error["loc"])
            error_details.append(f"  {field}: {error['msg']}")

        error_msg = f"Configuration validation failed for {path}:\n" + "\n".join(
            error_details
        )
        raise ValueError(error_msg) from e
