"""Model-kind detection and loading for the CLI and HTTP server."""

from __future__ import annotations

import json
import os

from .engine import PlutoEngine
from .ensemble import EnsembleEngine
from .task_specialist import TaskSpecialistEngine


def resolve_model_path(model: str, token: str | None = None) -> str:
    """Return a local path, downloading a Hub repository if needed."""
    if os.path.exists(model):
        return model
    from huggingface_hub import snapshot_download

    return snapshot_download(model, token=token)


def _ensemble_config(directory: str) -> str | None:
    preferred = os.path.join(directory, "ensemble_config.json")
    if os.path.exists(preferred):
        return preferred
    legacy = os.path.join(directory, "config.json")
    if os.path.exists(legacy):
        try:
            data = json.load(open(legacy, encoding="utf-8"))
        except Exception:
            return None
        if isinstance(data, dict) and "models" in data and "weights" in data:
            return legacy
    return None


def load_engine(model: str, device: str | None = None, token: str | None = None):
    """Load an ensemble, specialist, or general checkpoint by its layout."""
    path = resolve_model_path(model, token=token)
    if os.path.isdir(path):
        ensemble = _ensemble_config(path)
        if ensemble:
            return EnsembleEngine(ensemble, device=device, token=token)
        if os.path.exists(os.path.join(path, "pluto_task_config.json")):
            return TaskSpecialistEngine(path, device=device)
        if os.path.exists(os.path.join(path, "pluto_config.json")):
            return PlutoEngine(path, device=device)
        raise ValueError(
            f"{path!r} is not a Pluto checkpoint: expected ensemble_config.json, "
            "pluto_task_config.json, or pluto_config.json"
        )
    if os.path.isfile(path):
        base = os.path.dirname(path)
        name = os.path.basename(path)
        if name in ("ensemble_config.json", "config.json"):
            ensemble = _ensemble_config(base)
            if ensemble:
                return EnsembleEngine(ensemble, device=device, token=token)
        if name == "pluto_task_config.json":
            return TaskSpecialistEngine(base, device=device)
        if name == "pluto_config.json":
            return PlutoEngine(base, device=device)
    raise ValueError(f"unsupported model path: {path!r}")