cache_manager.py

"""
Cache Manager — scansione cartelle locali con modelli scaricati/convertiti.

Riconosce:
  ir_complete   → .xml + .bin accoppiati, size > 0, config.json presente
  ir_partial    → .xml senza .bin corrispondente, o file 0-byte
  raw_complete  → config.json + ≥1 .safetensors/.bin/.onnx, nessun .xml
  raw_partial   → config.json assente, o file .incomplete presenti
  unknown       → nulla di riconoscibile

Azioni disponibili per stato:
  ir_complete   → test_inference | open_folder | delete
  ir_partial    → resume_download | delete
  raw_complete  → convert_local | open_folder | delete
  raw_partial   → resume_download | delete
  unknown       → open_folder | delete
"""

import json
import logging
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Optional

log = logging.getLogger(__name__)


# ── Tipi ──────────────────────────────────────────────────────────────────────
class ModelState(str, Enum):
    IR_COMPLETE  = "ir_complete"
    IR_PARTIAL   = "ir_partial"
    RAW_COMPLETE = "raw_complete"
    RAW_PARTIAL  = "raw_partial"
    UNKNOWN      = "unknown"


class Action(str, Enum):
    TEST_INFERENCE  = "test_inference"
    CONVERT_LOCAL   = "convert_local"
    RESUME_DOWNLOAD = "resume_download"
    OPEN_FOLDER     = "open_folder"
    DELETE          = "delete"


_STATE_ACTIONS: dict[ModelState, list[Action]] = {
    ModelState.IR_COMPLETE:  [Action.TEST_INFERENCE, Action.OPEN_FOLDER, Action.DELETE],
    ModelState.IR_PARTIAL:   [Action.RESUME_DOWNLOAD, Action.OPEN_FOLDER, Action.DELETE],
    ModelState.RAW_COMPLETE: [Action.CONVERT_LOCAL, Action.OPEN_FOLDER, Action.DELETE],
    ModelState.RAW_PARTIAL:  [Action.RESUME_DOWNLOAD, Action.OPEN_FOLDER, Action.DELETE],
    ModelState.UNKNOWN:      [Action.OPEN_FOLDER, Action.DELETE],
}

_STATE_LABELS = {
    ModelState.IR_COMPLETE:  ("IR ✓ Completo",    "#4CAF50"),
    ModelState.IR_PARTIAL:   ("IR ⚠ Incompleto",  "#FF9800"),
    ModelState.RAW_COMPLETE: ("RAW ✓ Pronto",     "#5FB3FF"),
    ModelState.RAW_PARTIAL:  ("RAW ⚠ Incompleto", "#FF9800"),
    ModelState.UNKNOWN:      ("? Sconosciuto",     "#555"),
}


@dataclass
class FileInfo:
    name: str
    size_bytes: int
    is_complete: bool   # False se 0-byte o .incomplete


@dataclass
class CachedModel:
    path: Path
    state: ModelState
    repo_id: Optional[str]        # da config.json se leggibile
    architecture: Optional[str]   # da config.json
    category: str                 # llm_chat | vlm | stt | embedding | diffusion | unknown
    size_gb: float
    files: list[FileInfo] = field(default_factory=list)
    issues: list[str] = field(default_factory=list)   # problemi di integrità trovati

    @property
    def state_label(self) -> str:
        return _STATE_LABELS[self.state][0]

    @property
    def state_color(self) -> str:
        return _STATE_LABELS[self.state][1]

    @property
    def available_actions(self) -> list[Action]:
        return _STATE_ACTIONS[self.state]

    @property
    def display_name(self) -> str:
        if self.repo_id:
            return self.repo_id.split("/")[-1]
        return self.path.name


# ── Helpers interni ───────────────────────────────────────────────────────────
def _read_config(path: Path) -> dict:
    """Legge config.json dalla cartella, ritorna {} su errore."""
    cfg_path = path / "config.json"
    if not cfg_path.exists():
        return {}
    try:
        with open(cfg_path, encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return {}


def _guess_category(cfg: dict) -> str:
    arch = (cfg.get("architectures") or [""])[0].lower()
    model_type = cfg.get("model_type", "").lower()
    combined = arch + " " + model_type
    if any(x in combined for x in ["whisper", "speech", "wav2vec", "hubert"]):
        return "stt"
    if any(x in combined for x in ["vl", "llava", "vision", "blip", "qwen2_5_vl",
                                     "qwen3_5", "internvl", "minicpm"]):
        return "vlm"
    if any(x in combined for x in ["bert", "roberta", "embedding", "sentence"]):
        return "embedding"
    if any(x in combined for x in ["diffusion", "unet", "vae", "stable"]):
        return "diffusion"
    if arch or model_type:
        return "llm_chat"
    return "unknown"


def _dir_size_gb(path: Path) -> float:
    """Dimensione totale della cartella in GB, escludendo .cache/ di HF hub."""
    total = 0
    for f in path.rglob("*"):
        if not f.is_file():
            continue
        parts = f.relative_to(path).parts
        if ".cache" in parts:
            continue
        total += f.stat().st_size
    return round(total / (1024 ** 3), 3)


def _collect_files(path: Path) -> list[FileInfo]:
    """Lista file diretti e in sotto-cartelle, con flag integrità.
    Esclude la cartella .cache/huggingface/download/ (lock file HF hub).
    """
    infos = []
    for f in sorted(path.rglob("*")):
        if not f.is_file():
            continue
        # salta i lock file di hf_hub_download — non sono file modello
        rel = f.relative_to(path)
        parts = rel.parts
        if ".cache" in parts or f.name.endswith(".lock"):
            continue
        size = f.stat().st_size
        is_complete = (size > 0) and not f.name.endswith(".incomplete")
        infos.append(FileInfo(
            name=str(rel),
            size_bytes=size,
            is_complete=is_complete,
        ))
    return infos


# ── Classificatore principale ─────────────────────────────────────────────────
def classify_model_dir(path: Path) -> CachedModel:
    """
    Analizza una singola cartella e restituisce un CachedModel con stato e problemi.
    """
    cfg      = _read_config(path)
    repo_id  = cfg.get("_name_or_path") or cfg.get("name_or_path") or None
    arch_list = cfg.get("architectures") or []
    arch     = arch_list[0] if arch_list else None
    category = _guess_category(cfg)
    files    = _collect_files(path)
    issues   = []

    # file .incomplete → download interrotto
    incomplete = [f for f in files if not f.is_complete]
    if incomplete:
        for fi in incomplete:
            issues.append(f"File incompleto: {fi.name}")

    # ── Classifica per contenuto ──────────────────────────────────────────────
    xml_files  = [f for f in files if f.name.endswith(".xml")]
    bin_files  = [f for f in files if f.name.endswith(".bin")]
    sf_files   = [f for f in files if f.name.endswith(".safetensors")]
    onnx_files = [f for f in files if f.name.endswith(".onnx")]
    has_config = (path / "config.json").exists()

    # ── IR: presenza di .xml ──────────────────────────────────────────────────
    if xml_files:
        # verifica accoppiamento xml↔bin
        xml_stems = {Path(f.name).stem for f in xml_files}
        bin_stems = {Path(f.name).stem for f in bin_files}
        orphan_xml = xml_stems - bin_stems
        orphan_bin = bin_stems - xml_stems

        zero_byte = [f for f in xml_files + bin_files if f.size_bytes == 0]

        for stem in orphan_xml:
            issues.append(f"XML senza BIN: {stem}.xml")
        for stem in orphan_bin:
            issues.append(f"BIN senza XML: {stem}.bin")
        for f in zero_byte:
            issues.append(f"File 0 byte: {f.name}")

        if issues or incomplete:
            state = ModelState.IR_PARTIAL
        else:
            state = ModelState.IR_COMPLETE

        return CachedModel(
            path=path, state=state, repo_id=repo_id,
            architecture=arch, category=category,
            size_gb=_dir_size_gb(path), files=files, issues=issues,
        )

    # ── RAW: safetensors / onnx / pytorch bin ─────────────────────────────────
    raw_weights = sf_files + onnx_files + [
        f for f in bin_files if not f.name.startswith("openvino")
    ]
    if raw_weights:
        if not has_config:
            issues.append("config.json assente")
        bad = [f for f in raw_weights if not f.is_complete]
        for f in bad:
            issues.append(f"Peso incompleto: {f.name}")

        if issues or incomplete:
            state = ModelState.RAW_PARTIAL
        else:
            state = ModelState.RAW_COMPLETE

        return CachedModel(
            path=path, state=state, repo_id=repo_id,
            architecture=arch, category=category,
            size_gb=_dir_size_gb(path), files=files, issues=issues,
        )

    # ── Niente di riconoscibile ───────────────────────────────────────────────
    if not files:
        issues.append("Cartella vuota")
    else:
        issues.append(f"Nessun file modello riconoscibile ({len(files)} file presenti)")

    return CachedModel(
        path=path, state=ModelState.UNKNOWN, repo_id=repo_id,
        architecture=arch, category=category,
        size_gb=_dir_size_gb(path), files=files, issues=issues,
    )


# ── Scanner cartella ──────────────────────────────────────────────────────────
def scan_cache_dir(base_dir: Path, max_depth: int = 2) -> list[CachedModel]:
    """
    Scansiona base_dir cercando cartelle che contengono modelli.
    Esplora fino a max_depth livelli. Restituisce lista ordinata per path.

    Strategia: una cartella è un candidato se contiene almeno uno di:
    config.json, *.xml, *.safetensors, *.onnx, *.bin, *.incomplete
    Non esplora sotto-cartelle di un candidato già trovato.
    """
    if not base_dir.exists() or not base_dir.is_dir():
        log.warning(f"scan_cache_dir: {base_dir} non esiste o non è una dir")
        return []

    _MODEL_MARKERS = {
        "config.json", ".xml", ".safetensors", ".onnx", ".bin", ".incomplete"
    }

    def _is_model_dir(path: Path) -> bool:
        try:
            for child in path.iterdir():
                if not child.is_file():
                    continue
                if child.name in _MODEL_MARKERS:
                    return True
                if any(child.name.endswith(ext) for ext in _MODEL_MARKERS if ext.startswith(".")):
                    return True
        except PermissionError:
            pass
        return False

    candidates: list[Path] = []

    def _walk(current: Path, depth: int):
        if depth > max_depth:
            return
        try:
            subdirs = [c for c in current.iterdir() if c.is_dir()
                       and not c.name.startswith(".") and c.name != "__pycache__"]
        except PermissionError:
            return

        for subdir in subdirs:
            if _is_model_dir(subdir):
                candidates.append(subdir)
                # non scende dentro — i pesi sono file flat nella dir modello
            else:
                _walk(subdir, depth + 1)

    _walk(base_dir, 0)

    results = []
    for path in sorted(candidates):
        try:
            results.append(classify_model_dir(path))
        except Exception as e:
            log.warning(f"classify_model_dir fallita per {path}: {e}")

    log.info(f"scan_cache_dir({base_dir}): {len(results)} modelli trovati")
    return results