path: root/src/mcp_sd15/config.py

"""TOML configuration loading and model path confinement."""

from __future__ import annotations

import math
import tomllib
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .errors import ConfigError

VALID_SIZES = frozenset({"square_512", "landscape_512", "portrait_512"})
VALID_SCHEDULERS = frozenset({"dpmpp_2m_karras", "dpmpp_2m_sde_karras", "euler_a", "ddim", "lcm"})
VALID_TRANSPORTS = frozenset({"stdio", "streamable-http"})


@dataclass(frozen=True)
class LoRAConfig:
    file: str
    weight: float = 1.0


@dataclass(frozen=True)
class ModelConfig:
    root: Path
    checkpoint: str
    vae: str | None
    loras: tuple[LoRAConfig, ...]

    def checkpoint_path(self) -> Path:
        return self.root / "checkpoints" / self.checkpoint

    def vae_path(self) -> Path | None:
        return self.root / "vae" / self.vae if self.vae else None

    def lora_paths(self) -> tuple[Path, ...]:
        return tuple(self.root / "lora" / item.file for item in self.loras)


@dataclass(frozen=True)
class GenerationConfig:
    steps: int = 25
    cfg_scale: float = 7.0
    scheduler: str = "dpmpp_2m_karras"
    default_size: str = "square_512"


@dataclass(frozen=True)
class ControlModelConfig:
    enabled: bool
    model: str | None
    preprocessor: str | None = None


@dataclass(frozen=True)
class ControlNetConfig:
    depth: ControlModelConfig
    openpose: ControlModelConfig
    inpaint: ControlModelConfig

    def for_mode(self, mode: str) -> ControlModelConfig:
        key = "openpose" if mode == "openpose_full" else mode
        if key not in {"depth", "openpose", "inpaint"}:
            raise ConfigError(f"unknown ControlNet mode: {mode}")
        return getattr(self, key)


@dataclass(frozen=True)
class OutputConfig:
    directory: Path
    enabled: bool = False


@dataclass(frozen=True)
class ServerConfig:
    transport: str = "stdio"
    host: str = "127.0.0.1"
    port: int = 8000


@dataclass(frozen=True)
class PromptConfig:
    positive: str = ""
    negative: str = ""


@dataclass(frozen=True)
class AppConfig:
    model: ModelConfig
    generation: GenerationConfig
    controlnet: ControlNetConfig
    output: OutputConfig
    server: ServerConfig = ServerConfig()
    prompt: PromptConfig = PromptConfig()

    def controlnet_path(self, mode: str) -> Path:
        item = self.controlnet.for_mode(mode)
        if not item.enabled or not item.model:
            raise ConfigError(f"ControlNet mode is disabled: {mode}")
        return self.model.root / "controlnet" / item.model


def _table(data: dict[str, Any], key: str, *, required: bool = False) -> dict[str, Any]:
    value = data.get(key)
    if value is None and not required:
        return {}
    if not isinstance(value, dict):
        raise ConfigError(f"[{key}] must be a TOML table")
    return value


def _string(table: dict[str, Any], key: str, *, required: bool = False) -> str | None:
    value = table.get(key)
    if value is None and not required:
        return None
    if not isinstance(value, str) or not value.strip():
        raise ConfigError(f"{key} must be a non-empty string")
    return value.strip()


def _filename(value: str | None, label: str) -> str | None:
    if value is None:
        return None
    path = Path(value)
    if path.name != value or path.is_absolute() or value in {".", ".."}:
        raise ConfigError(f"{label} must be a file name, not a path")
    return value


def _read_prompt(path: Path) -> str:
    try:
        return path.read_text(encoding="utf-8").strip()
    except OSError as exc:
        raise ConfigError(f"cannot read prompt file {path}: {exc}") from exc


def _bool(table: dict[str, Any], key: str, default: bool) -> bool:
    value = table.get(key, default)
    if not isinstance(value, bool):
        raise ConfigError(f"{key} must be a boolean")
    return value


def _number(table: dict[str, Any], key: str, default: float) -> float:
    value = table.get(key, default)
    if isinstance(value, bool) or not isinstance(value, int | float):
        raise ConfigError(f"{key} must be a number")
    result = float(value)
    if not math.isfinite(result):
        raise ConfigError(f"{key} must be finite")
    return result


def _control(table: dict[str, Any], name: str) -> ControlModelConfig:
    enabled = _bool(table, "enabled", False)
    model = _filename(_string(table, "model"), f"controlnet.{name}.model")
    if enabled and model is None:
        raise ConfigError(f"controlnet.{name}.model is required when enabled")
    preprocessor = _string(table, "preprocessor")
    if name == "depth" and preprocessor not in {None, "midas"}:
        raise ConfigError("controlnet.depth.preprocessor must be 'midas'")
    return ControlModelConfig(
        enabled, model, preprocessor or ("midas" if name == "depth" else None)
    )


def load_config(path: str | Path, *, validate_paths: bool = True) -> AppConfig:
    """Load a configuration file, validating all values and local model paths."""
    config_path = Path(path).expanduser().resolve()
    try:
        with config_path.open("rb") as stream:
            raw = tomllib.load(stream)
    except (OSError, tomllib.TOMLDecodeError) as exc:
        raise ConfigError(f"cannot read config {config_path}: {exc}") from exc

    model_raw = _table(raw, "model", required=True)
    root_value = _string(model_raw, "root", required=True)
    assert root_value is not None
    root = Path(root_value).expanduser()
    if not root.is_absolute():
        root = config_path.parent / root
    root = root.resolve()
    checkpoint = _filename(_string(model_raw, "checkpoint", required=True), "model.checkpoint")
    assert checkpoint is not None
    vae = _filename(_string(model_raw, "vae"), "model.vae")

    lora_raw = model_raw.get("lora", [])
    if not isinstance(lora_raw, list):
        raise ConfigError("model.lora must be an array of tables")
    loras: list[LoRAConfig] = []
    for index, value in enumerate(lora_raw):
        if not isinstance(value, dict):
            raise ConfigError(f"model.lora[{index}] must be a table")
        file = _filename(_string(value, "file", required=True), f"model.lora[{index}].file")
        weight = _number(value, "weight", 1.0)
        if not -4.0 <= weight <= 4.0:
            raise ConfigError(f"model.lora[{index}].weight must be between -4 and 4")
        assert file is not None
        loras.append(LoRAConfig(file, weight))
    if len({item.file for item in loras}) != len(loras):
        raise ConfigError("LoRA file names must be unique")

    generation_raw = _table(raw, "generation")
    steps_value = generation_raw.get("steps", 25)
    if (
        isinstance(steps_value, bool)
        or not isinstance(steps_value, int)
        or not 1 <= steps_value <= 150
    ):
        raise ConfigError("generation.steps must be an integer between 1 and 150")
    cfg_scale = _number(generation_raw, "cfg_scale", 7.0)
    if not 0 <= cfg_scale <= 30:
        raise ConfigError("generation.cfg_scale must be between 0 and 30")
    scheduler = _string(generation_raw, "scheduler") or "dpmpp_2m_karras"
    if scheduler not in VALID_SCHEDULERS:
        raise ConfigError(f"unsupported scheduler: {scheduler}")
    default_size = _string(generation_raw, "default_size") or "square_512"
    if default_size not in VALID_SIZES:
        raise ConfigError(f"unsupported default_size: {default_size}")

    control_raw = _table(raw, "controlnet")
    controlnet = ControlNetConfig(
        depth=_control(_table(control_raw, "depth"), "depth"),
        openpose=_control(_table(control_raw, "openpose"), "openpose"),
        inpaint=_control(_table(control_raw, "inpaint"), "inpaint"),
    )
    output_raw = _table(raw, "output")
    output_value = _string(output_raw, "directory") or "./outputs"
    output_dir = Path(output_value).expanduser()
    if not output_dir.is_absolute():
        output_dir = config_path.parent / output_dir

    server_raw = _table(raw, "server")
    transport = _string(server_raw, "transport") or "stdio"
    if transport not in VALID_TRANSPORTS:
        raise ConfigError(f"unsupported transport: {transport}")
    host = _string(server_raw, "host") or "127.0.0.1"
    port_value = server_raw.get("port", 8000)
    if (
        isinstance(port_value, bool)
        or not isinstance(port_value, int)
        or not 1 <= port_value <= 65535
    ):
        raise ConfigError("server.port must be an integer between 1 and 65535")

    prompt_raw = _table(raw, "prompt")
    positive_file = _filename(_string(prompt_raw, "positive_file"), "prompt.positive_file")
    negative_file = _filename(_string(prompt_raw, "negative_file"), "prompt.negative_file")
    prompt_positive = ""
    prompt_negative = ""
    if positive_file:
        positive_path = root / "prompts" / positive_file
        prompt_positive = _read_prompt(positive_path)
    if negative_file:
        negative_path = root / "prompts" / negative_file
        prompt_negative = _read_prompt(negative_path)

    config = AppConfig(
        model=ModelConfig(root, checkpoint, vae, tuple(loras)),
        generation=GenerationConfig(steps_value, cfg_scale, scheduler, default_size),
        controlnet=controlnet,
        output=OutputConfig(output_dir.resolve(), _bool(output_raw, "enabled", False)),
        server=ServerConfig(transport, host, port_value),
        prompt=PromptConfig(prompt_positive, prompt_negative),
    )
    if validate_paths:
        validate_model_paths(config)
    return config


def validate_model_paths(config: AppConfig) -> None:
    """Fail before GPU initialization when a configured model file is absent."""
    if not config.model.root.is_dir():
        raise ConfigError(f"model root directory not found: {config.model.root}")
    required = [config.model.checkpoint_path()]
    vae = config.model.vae_path()
    if vae:
        required.append(vae)
    required.extend(config.model.lora_paths())
    for mode in ("depth", "openpose", "inpaint"):
        item = config.controlnet.for_mode(mode)
        if item.enabled:
            required.append(config.controlnet_path(mode))
    missing = [str(path) for path in required if not path.is_file()]
    if missing:
        raise ConfigError("model file(s) not found: " + ", ".join(missing))