"""Paths, model-type mapping, and folder bootstrap for the Model Manager."""
from __future__ import annotations

import os
from pathlib import Path

# Root of the host `models/` directory (mounted read-write into this container).
MODELS_DIR = Path(os.environ.get("MODELS_DIR", "/models")).resolve()

# Persistent data dir for the SQLite database.
DATA_DIR = Path(os.environ.get("DATA_DIR", "/data")).resolve()

# Generated images directory (ComfyUI output), mounted read-write for the gallery.
OUTPUT_DIR = Path(os.environ.get("OUTPUT_DIR", "/output")).resolve()

DB_PATH = DATA_DIR / "manager.db"

# Ports of the sibling services, used by the device-routing landing page.
COMFYUI_PORT = os.environ.get("COMFYUI_PORT", "8188")
COMFYUIMINI_PORT = os.environ.get("COMFYUIMINI_PORT", "3000")

# Internal URL of the ComfyUI container (reachable over the docker network),
# used to proxy the "free memory / unload models" action.
COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://comfyui:8188")

# Image file types shown in the gallery.
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp"}

# StabilityMatrix-style model types mapped to ComfyUI's standard `models/` subfolders.
# `key` is the stable id used by the API/UI; `folder` is the on-disk subdirectory.
MODEL_TYPES: list[dict[str, str]] = [
    {"key": "checkpoint", "label": "Checkpoint", "folder": "checkpoints"},
    {"key": "lora", "label": "LoRA / LyCORIS", "folder": "loras"},
    {"key": "vae", "label": "VAE", "folder": "vae"},
    {"key": "embedding", "label": "Textual Inversion", "folder": "embeddings"},
    {"key": "controlnet", "label": "ControlNet", "folder": "controlnet"},
    {"key": "upscaler", "label": "Upscaler", "folder": "upscale_models"},
    {"key": "clip", "label": "CLIP", "folder": "clip"},
    {"key": "clip_vision", "label": "CLIP Vision", "folder": "clip_vision"},
    {"key": "text_encoder", "label": "Text Encoder", "folder": "text_encoders"},
    {"key": "diffusion_model", "label": "Diffusion Model (UNET)", "folder": "diffusion_models"},
    {"key": "unet", "label": "UNET", "folder": "unet"},
    {"key": "hypernetwork", "label": "Hypernetwork", "folder": "hypernetworks"},
    {"key": "style_model", "label": "Style Model", "folder": "style_models"},
    {"key": "gligen", "label": "GLIGEN", "folder": "gligen"},
    {"key": "vae_approx", "label": "VAE Approx", "folder": "vae_approx"},
    {"key": "ipadapter", "label": "IP-Adapter", "folder": "ipadapter"},
    {"key": "other", "label": "Other", "folder": "other"},
]

# Quick lookups.
TYPE_BY_KEY: dict[str, dict[str, str]] = {t["key"]: t for t in MODEL_TYPES}
FOLDER_BY_KEY: dict[str, str] = {t["key"]: t["folder"] for t in MODEL_TYPES}
KEY_BY_FOLDER: dict[str, str] = {t["folder"]: t["key"] for t in MODEL_TYPES}


def folder_for_type(type_key: str) -> str:
    """Return the on-disk subfolder for a model type key, defaulting to `other`."""
    return FOLDER_BY_KEY.get(type_key, "other")


def ensure_dirs() -> None:
    """Create the data dir and all standard model subfolders if missing."""
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    for t in MODEL_TYPES:
        (MODELS_DIR / t["folder"]).mkdir(parents=True, exist_ok=True)


def safe_output_path(rel_path: str) -> Path:
    """Resolve a gallery-relative path to an absolute one under OUTPUT_DIR.

    Raises ValueError on any attempt to escape the output directory.
    """
    # Normalize and strip leading separators so it's treated as relative.
    rel = rel_path.replace("\\", "/").lstrip("/")
    target = (OUTPUT_DIR / rel).resolve()
    if target != OUTPUT_DIR and not str(target).startswith(str(OUTPUT_DIR) + os.sep):
        raise ValueError("Path is outside the output directory")
    return target