"""Discover the YOLO models available for detection.

ADetailer-style layout: drop weights under ``models/yolo/<category>/*.pt``. The
*category* (the sub-folder) becomes the detection label and its overlay colour, so
e.g. ``models/yolo/mosaic/lada.pt`` tags its boxes "mosaic" and ``models/yolo/face/
yolov8n-face.pt`` tags "face". The user ticks which discovered models are active; a
detect runs every ticked model and merges the results (see ``multi.MultiYoloDetector``).
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

YOLO_SUBDIR = ("models", "yolo")   # relative to the working directory
_UNCATEGORIZED = "misc"            # category for a .pt sitting directly under models/yolo


@dataclass(frozen=True)
class ModelEntry:
    path: str       # absolute path to the .pt
    category: str   # sub-folder under models/yolo (the detection label / colour key)
    name: str       # filename stem (display name)


def yolo_root(root: Path | None = None) -> Path:
    return (root or Path.cwd()).joinpath(*YOLO_SUBDIR)


def _category_of(pt: Path, root: Path) -> str:
    rel = pt.relative_to(root).parts
    return rel[0] if len(rel) > 1 else _UNCATEGORIZED


def discover_models(root: Path | None = None) -> list[ModelEntry]:
    """All ``*.pt`` under ``models/yolo/**``, sorted by (category, name)."""
    base = yolo_root(root)
    if not base.is_dir():
        return []
    out = [
        ModelEntry(path=str(p), category=_category_of(p, base), name=p.stem)
        for p in base.rglob("*.pt")
    ]
    out.sort(key=lambda e: (e.category.lower(), e.name.lower()))
    return out


def category_of(model_path: str, root: Path | None = None) -> str:
    """Category (label) for a model path, derived from its folder under models/yolo."""
    base = yolo_root(root)
    p = Path(model_path)
    try:
        return _category_of(p, base)
    except ValueError:  # outside models/yolo — fall back to the parent folder name
        return p.parent.name or _UNCATEGORIZED