"""
Model type detector.

Inspects a locally installed model directory to determine which pipeline
category it belongs to, without loading the model weights.

Detection order:
  1. model_index.json present  → diffusers pipeline → IMAGE_GENERATION
  2. config.json pipeline_tag  → mapped to a category
  3. Fallback                  → TEXT_GENERATION
"""

from __future__ import annotations

import json
from pathlib import Path

# ── categories ────────────────────────────────────────────────────────────────

TEXT_GENERATION     = "text-generation"
IMAGE_GENERATION    = "image-generation"
VISION              = "vision"
AUDIO_TRANSCRIPTION = "audio-transcription"
AUDIO_SPEECH        = "audio-speech"

# ── pipeline_tag → category ───────────────────────────────────────────────────

_TAG_MAP: dict[str, str] = {
    "text-generation":               TEXT_GENERATION,
    "text2text-generation":          TEXT_GENERATION,
    "conversational":                TEXT_GENERATION,
    "text-to-image":                 IMAGE_GENERATION,
    "image-to-image":                IMAGE_GENERATION,
    "inpainting":                    IMAGE_GENERATION,
    "image-to-text":                 VISION,
    "visual-question-answering":     VISION,
    "image-text-to-text":            VISION,
    "automatic-speech-recognition":  AUDIO_TRANSCRIPTION,
    "audio-classification":          AUDIO_TRANSCRIPTION,
    "text-to-speech":                AUDIO_SPEECH,
    "text-to-audio":                 AUDIO_SPEECH,
}

# ── model_type → category (config.json "model_type" field) ───────────────────
# pipeline_tag is from HF model card YAML and is often absent in local
# config.json.  model_type is always present and is a reliable fallback.

_MODEL_TYPE_MAP: dict[str, str] = {
    # image generation / diffusion
    "uvit":             IMAGE_GENERATION,   # amused
    "vqgan":            IMAGE_GENERATION,
    "ddpm":             IMAGE_GENERATION,
    "consistency":      IMAGE_GENERATION,
    # vision (image understanding)
    "clip":             VISION,
    "blip":             VISION,
    "blip-2":           VISION,
    "idefics":          VISION,
    "llava":            VISION,
    "pix2struct":       VISION,
    "florence":         VISION,
    "siglip":           VISION,
    # audio transcription (ASR)
    "whisper":          AUDIO_TRANSCRIPTION,
    "wav2vec2":         AUDIO_TRANSCRIPTION,
    "hubert":           AUDIO_TRANSCRIPTION,
    "data2vec-audio":   AUDIO_TRANSCRIPTION,
    "unispeech":        AUDIO_TRANSCRIPTION,
    "sew":              AUDIO_TRANSCRIPTION,
    "wavlm":            AUDIO_TRANSCRIPTION,
    # audio speech (TTS)
    "speecht5":         AUDIO_SPEECH,
    "vits":             AUDIO_SPEECH,
    "encodec":          AUDIO_SPEECH,
}

# ── public API ────────────────────────────────────────────────────────────────

def detect(model_dir: Path) -> str:
    """
    Return the pipeline category for the HF model at model_dir.

    Returns one of the category constants defined in this module
    (TEXT_GENERATION, IMAGE_GENERATION, VISION, AUDIO_TRANSCRIPTION,
    AUDIO_SPEECH).
    """
    # 1. Diffusers pipelines always ship a model_index.json
    if (model_dir / "model_index.json").exists():
        return IMAGE_GENERATION

    # 2. Standard HF config.json — check pipeline_tag then model_type
    config_path = model_dir / "config.json"
    if config_path.exists():
        try:
            cfg = json.loads(config_path.read_text(encoding="utf-8"))

            tag = cfg.get("pipeline_tag", "")
            if tag in _TAG_MAP:
                return _TAG_MAP[tag]

            model_type = cfg.get("model_type", "").lower()
            if model_type in _MODEL_TYPE_MAP:
                return _MODEL_TYPE_MAP[model_type]

        except (json.JSONDecodeError, OSError):
            pass

    # 3. Default — assume causal text generation
    return TEXT_GENERATION


def label(category: str) -> str:
    """Return a human-readable label for a category constant."""
    return {
        TEXT_GENERATION:     "Text Generation",
        IMAGE_GENERATION:    "Image Generation (Diffusion)",
        VISION:              "Vision (Image-to-Text)",
        AUDIO_TRANSCRIPTION: "Audio Transcription (ASR)",
        AUDIO_SPEECH:        "Audio Speech (TTS)",
    }.get(category, category)
