cubbi/mcontainer/config.py

import yaml
from pathlib import Path
from typing import Dict, Optional

from .models import Config, Driver

DEFAULT_CONFIG_DIR = Path.home() / ".config" / "mc"
DEFAULT_CONFIG_FILE = DEFAULT_CONFIG_DIR / "config.yaml"
DEFAULT_DRIVERS_DIR = Path.home() / ".config" / "mc" / "drivers"
PROJECT_ROOT = Path(__file__).parent.parent
BUILTIN_DRIVERS_DIR = Path(__file__).parent / "drivers"  # mcontainer/drivers

# Dynamically loaded from drivers directory at runtime
DEFAULT_DRIVERS = {}


class ConfigManager:
    def __init__(self, config_path: Optional[Path] = None):
        self.config_path = config_path or DEFAULT_CONFIG_FILE
        self.config_dir = self.config_path.parent
        self.drivers_dir = DEFAULT_DRIVERS_DIR
        self.config = self._load_or_create_config()

        # Always load package drivers on initialization
        # These are separate from the user config
        self.builtin_drivers = self._load_package_drivers()

    def _load_or_create_config(self) -> Config:
        """Load existing config or create a new one with defaults"""
        if self.config_path.exists():
            try:
                with open(self.config_path, "r") as f:
                    config_data = yaml.safe_load(f) or {}

                # Create a new config from scratch, then update with data from file
                config = Config(
                    docker=config_data.get("docker", {}),
                    defaults=config_data.get("defaults", {}),
                )

                # Add drivers
                if "drivers" in config_data:
                    for driver_name, driver_data in config_data["drivers"].items():
                        config.drivers[driver_name] = Driver.model_validate(driver_data)

                return config
            except Exception as e:
                print(f"Error loading config: {e}")
                return self._create_default_config()
        else:
            return self._create_default_config()

    def _create_default_config(self) -> Config:
        """Create a default configuration"""
        self.config_dir.mkdir(parents=True, exist_ok=True)
        self.drivers_dir.mkdir(parents=True, exist_ok=True)

        # Initial config without drivers
        config = Config(
            docker={
                "socket": "/var/run/docker.sock",
                "network": "mc-network",
            },
            defaults={
                "driver": "goose",
            },
        )

        self.save_config(config)
        return config

    def save_config(self, config: Optional[Config] = None) -> None:
        """Save the current config to disk"""
        if config:
            self.config = config

        self.config_dir.mkdir(parents=True, exist_ok=True)

        # Use model_dump with mode="json" for proper serialization of enums
        config_dict = self.config.model_dump(mode="json")

        # Write to file
        with open(self.config_path, "w") as f:
            yaml.dump(config_dict, f)

    def get_driver(self, name: str) -> Optional[Driver]:
        """Get a driver by name, checking builtin drivers first, then user-configured ones"""
        # Check builtin drivers first (package drivers take precedence)
        if name in self.builtin_drivers:
            return self.builtin_drivers[name]
        # If not found, check user-configured drivers
        return self.config.drivers.get(name)

    def list_drivers(self) -> Dict[str, Driver]:
        """List all available drivers (both builtin and user-configured)"""
        # Start with user config drivers
        all_drivers = dict(self.config.drivers)

        # Add builtin drivers, overriding any user drivers with the same name
        # This ensures that package-provided drivers always take precedence
        all_drivers.update(self.builtin_drivers)

        return all_drivers

    # Session management has been moved to SessionManager in session.py

    def load_driver_from_dir(self, driver_dir: Path) -> Optional[Driver]:
        """Load a driver configuration from a directory"""
        # Try with mc-driver.yaml first (new format), then mai-driver.yaml (legacy)
        yaml_path = driver_dir / "mc-driver.yaml"
        if not yaml_path.exists():
            yaml_path = driver_dir / "mai-driver.yaml"  # Backward compatibility
            if not yaml_path.exists():
                return None

        try:
            with open(yaml_path, "r") as f:
                driver_data = yaml.safe_load(f)

            # Extract required fields
            if not all(
                k in driver_data
                for k in ["name", "description", "version", "maintainer"]
            ):
                print(f"Driver config {yaml_path} missing required fields")
                return None

            # Use Driver.model_validate to handle all fields from YAML
            # This will map all fields according to the Driver model structure
            try:
                # Ensure image field is set if not in YAML
                if "image" not in driver_data:
                    driver_data["image"] = f"monadical/mc-{driver_data['name']}:latest"

                driver = Driver.model_validate(driver_data)
                return driver
            except Exception as validation_error:
                print(
                    f"Error validating driver data from {yaml_path}: {validation_error}"
                )
                return None

        except Exception as e:
            print(f"Error loading driver from {yaml_path}: {e}")
            return None

    def _load_package_drivers(self) -> Dict[str, Driver]:
        """Load all package drivers from the mcontainer/drivers directory"""
        drivers = {}

        if not BUILTIN_DRIVERS_DIR.exists():
            return drivers

        # Search for mc-driver.yaml files in each subdirectory
        for driver_dir in BUILTIN_DRIVERS_DIR.iterdir():
            if driver_dir.is_dir():
                driver = self.load_driver_from_dir(driver_dir)
                if driver:
                    drivers[driver.name] = driver

        return drivers

    def get_driver_path(self, driver_name: str) -> Optional[Path]:
        """Get the directory path for a driver"""
        # Check package drivers first (these are the bundled ones)
        package_path = BUILTIN_DRIVERS_DIR / driver_name
        if package_path.exists() and package_path.is_dir():
            return package_path

        # Then check user drivers
        user_path = self.drivers_dir / driver_name
        if user_path.exists() and user_path.is_dir():
            return user_path

        return None