"""Configuration file support for ORBIT (~/.orbit/config.yaml)."""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


CONFIG_DIR = Path.home() / ".orbit"
CONFIG_FILE = CONFIG_DIR / "config.yaml"


@dataclass
class OrbitConfig:
    """User configuration loaded from ~/.orbit/config.yaml."""

    default_policy: str = "auto"
    default_min_grade: str = "C"
    gpu_memory: int | None = None
    google_api_key: str | None = None
    default_framework: str = "lerobot"
    default_format: str | None = None
    default_device: str | None = None
    coach: dict = field(default_factory=lambda: {"target_episodes": 50, "poll_interval": 2.0})
    plugins_dir: str | None = None

    @classmethod
    def load(cls) -> "OrbitConfig":
        """Load config from ~/.orbit/config.yaml, falling back to defaults."""
        if not CONFIG_FILE.exists():
            return cls()

        try:
            import yaml
        except ImportError:
            # yaml not installed — silently use defaults
            return cls()

        try:
            with open(CONFIG_FILE) as f:
                data = yaml.safe_load(f) or {}
        except Exception:
            return cls()

        return cls._from_dict(data)

    @classmethod
    def _from_dict(cls, data: dict[str, Any]) -> "OrbitConfig":
        """Build config from a parsed YAML dict, expanding env vars."""
        cfg = cls()
        for key in [
            "default_policy", "default_min_grade", "gpu_memory",
            "google_api_key", "default_framework", "default_format",
            "default_device", "plugins_dir",
        ]:
            if key in data:
                val = data[key]
                if isinstance(val, str) and val.startswith("${") and val.endswith("}"):
                    val = os.environ.get(val[2:-1], None)
                setattr(cfg, key, val)
        if "coach" in data and isinstance(data["coach"], dict):
            cfg.coach = {**cfg.coach, **data["coach"]}
        return cfg

    def save(self) -> None:
        """Save current config to ~/.orbit/config.yaml."""
        try:
            import yaml
        except ImportError:
            raise RuntimeError("PyYAML required to save config: pip install pyyaml")

        CONFIG_DIR.mkdir(parents=True, exist_ok=True)
        data = {
            "default_policy": self.default_policy,
            "default_min_grade": self.default_min_grade,
            "default_framework": self.default_framework,
        }
        if self.gpu_memory is not None:
            data["gpu_memory"] = self.gpu_memory
        if self.google_api_key is not None:
            data["google_api_key"] = self.google_api_key
        if self.default_format is not None:
            data["default_format"] = self.default_format
        if self.default_device is not None:
            data["default_device"] = self.default_device
        if self.plugins_dir is not None:
            data["plugins_dir"] = self.plugins_dir
        if self.coach != {"target_episodes": 50, "poll_interval": 2.0}:
            data["coach"] = self.coach

        with open(CONFIG_FILE, "w") as f:
            yaml.dump(data, f, default_flow_style=False, sort_keys=False)


_VALID_VALUES: dict[str, list[str] | None] = {
    "default_policy": ["auto", "act", "diffusion_policy", "smolvla", "bc", "dp3"],
    "default_min_grade": ["A", "B", "C", "D"],
    "default_framework": ["lerobot"],
    "default_format": [None, "auto", "lerobot", "lerobot-local", "rlds", "hdf5", "rosbag", "directory"],
    "default_device": [None, "cpu", "cuda", "mps"],
}


def validate_config_value(key: str, value: Any) -> str | None:
    """Validate a config value before setting. Returns error message or None."""
    allowed = _VALID_VALUES.get(key)
    if allowed is not None and value not in allowed:
        options = [str(v) for v in allowed if v is not None]
        return f"Invalid value '{value}' for {key}. Allowed: {', '.join(options)}"

    if key == "gpu_memory" and value is not None:
        if not isinstance(value, (int, float)) or value <= 0:
            return f"gpu_memory must be a positive number, got {value}"

    return None


def get_config() -> OrbitConfig:
    """Get the global config (loads once, caches)."""
    if not hasattr(get_config, "_cached"):
        get_config._cached = OrbitConfig.load()
    return get_config._cached
