"""Domain configuration loader — loads built-in or custom YAML domain configs."""

from __future__ import annotations

import os
from pathlib import Path
from typing import Union

import yaml

from res_sum.schemas import DomainConfig, EntityTypeDef


# Registry of built-in domains (lazy-loaded)
_BUILTIN_DOMAINS: dict[str, DomainConfig] = {}


def _register_builtins() -> None:
    """Register built-in domain configs on first access."""
    if _BUILTIN_DOMAINS:
        return

    from res_sum.domains.ecology import ECOLOGY_CONFIG
    from res_sum.domains.general import GENERAL_CONFIG

    _BUILTIN_DOMAINS["ecology"] = ECOLOGY_CONFIG
    _BUILTIN_DOMAINS["general"] = GENERAL_CONFIG


def load_domain(domain: Union[str, Path, DomainConfig]) -> DomainConfig:
    """Load a domain configuration.

    Args:
        domain: One of:
            - A built-in domain name: "ecology", "general"
            - A path to a YAML config file: "./my_domain.yaml"
            - A DomainConfig object (returned as-is)

    Returns:
        DomainConfig instance.
    """
    if isinstance(domain, DomainConfig):
        return domain

    domain_str = str(domain)

    # Check built-in domains first
    _register_builtins()
    if domain_str in _BUILTIN_DOMAINS:
        return _BUILTIN_DOMAINS[domain_str]

    # Try loading as YAML file path
    path = Path(domain_str)
    if path.exists() and path.suffix in (".yaml", ".yml"):
        return _load_yaml_domain(path)

    available = ", ".join(list(_BUILTIN_DOMAINS.keys()))
    raise ValueError(
        f"Unknown domain '{domain_str}'. "
        f"Available built-in domains: {available}. "
        f"Or provide a path to a YAML config file."
    )


def _load_yaml_domain(path: Path) -> DomainConfig:
    """Load a DomainConfig from a YAML file."""
    with open(path, "r") as f:
        data = yaml.safe_load(f)

    entity_types = [
        EntityTypeDef(**et) for et in data.get("entity_types", [])
    ]

    return DomainConfig(
        name=data.get("name", path.stem),
        entity_types=entity_types,
        relationship_types=data.get("relationship_types", []),
        cot_prompt=data.get("cot_prompt", ""),
        section_headers=data.get("section_headers", []),
    )


def list_domains() -> list[str]:
    """Return names of available built-in domains."""
    _register_builtins()
    return list(_BUILTIN_DOMAINS.keys())
