conversion_engine.py

"""
Conversion engine — wrappa optimum-cli come subprocess.
Parsing stdout in real-time per estrarre progresso e messaggi.
Progettato per girare in QThread con signal emission.
"""

import re
import shutil
import subprocess
import tempfile
import textwrap
import logging
from pathlib import Path
from dataclasses import dataclass, field
from typing import Callable, Optional

log = logging.getLogger(__name__)

# ── Auto-discovery optimum-cli e python3 ─────────────────────────────────────
#
# Priorità ricerca optimum-cli:
#   1. Path configurato dall'utente in Settings (override esplicito)
#   2. Stessa bin/ del Python che gira OVForge (caso più comune: siamo nell'env)
#   3. PATH di sistema (fallback)
#
# Il path configurato viene passato via _set_env_override() da SettingsDialog
# dopo che l'utente lo ha salvato. Senza override, si usa l'auto-discovery.

import sys as _sys

_env_cli_override: Optional[str] = None   # impostato da set_env_override()
_env_py_override:  Optional[str] = None

def set_env_override(cli_path: Optional[str], py_path: Optional[str]):
    """
    Chiamato da main_window dopo il caricamento delle settings.
    Permette all'utente di sovrascrivere i path auto-discovered.
    """
    global _env_cli_override, _env_py_override
    _env_cli_override = cli_path or None
    _env_py_override  = py_path  or None
    log.info(f"env override: cli={cli_path!r}  py={py_path!r}")

def _find_optimum_cli() -> Optional[str]:
    """
    Trova optimum-cli nell'ordine: override utente → bin/ accanto a sys.executable → PATH.
    """
    if _env_cli_override and Path(_env_cli_override).exists():
        return _env_cli_override

    # bin/ dello stesso Python che sta girando (caso env attivato)
    candidate = Path(_sys.executable).parent / "optimum-cli"
    if candidate.exists():
        return str(candidate)

    # PATH sistema
    found = shutil.which("optimum-cli")
    if found:
        return found

    return None

def _find_python() -> str:
    """
    Trova il Python corretto: override utente → sys.executable → 'python3'.
    """
    if _env_py_override and Path(_env_py_override).exists():
        return _env_py_override

    # sys.executable è sempre valido (è il Python che sta girando OVForge)
    return _sys.executable

# ── Architetture con _text variant nativa in optimum-intel ───────────────────
# Per queste basta --task text-generation-with-past, optimum-intel fa il resto.
# Per le altre serve fn_get_submodels via script Python.
_NATIVE_TEXT_VARIANTS = {
    "qwen2_vl", "qwen2_5_vl", "qwen3_5", "qwen3_5_moe",
    "gemma3", "gemma4", "qwen3_vl",
}

def _needs_python_script(model_type: str) -> bool:
    """True se l'architettura NON ha _text variant nativa → serve script Python."""
    return model_type.lower() not in _NATIVE_TEXT_VARIANTS

# ── Strutture job ─────────────────────────────────────────────────────────────
@dataclass
class ConversionJob:
    job_id: str
    repo_id: str
    output_dir: Path
    cli_args: list[str]          # argv completo già costruito da profile_resolver
    estimated_size_gb: float = -1.0

@dataclass
class ConversionProgress:
    job_id: str
    percent: int                 # 0-100
    message: str
    phase: str                   # loading | exporting | quantizing | saving | done | error
    error: str = ""

# ── Pattern stdout optimum-cli ────────────────────────────────────────────────
# La CLI non ha progress JSON nativo — parseo l'output testuale
_PATTERNS = [
    # caricamento modello
    (re.compile(r"Loading model|Downloading model|Fetching", re.I), "loading", 5),
    (re.compile(r"config\.json|tokenizer|Loading checkpoint", re.I), "loading", 10),
    # esportazione
    (re.compile(r"Converting|Exporting|openvino_model", re.I), "exporting", 30),
    (re.compile(r"Traced|forward pass|Dynamic shape", re.I), "exporting", 50),
    # quantizzazione
    (re.compile(r"Compression|quantiz|NNCF|INT4|INT8|AWQ|scale estimation", re.I), "quantizing", 65),
    (re.compile(r"calibrat|wikitext|dataset", re.I), "quantizing", 70),
    (re.compile(r"Finished compression|Compression done", re.I), "quantizing", 85),
    # salvataggio
    (re.compile(r"Saving|Writing|openvino_model\.xml|openvino_model\.bin", re.I), "saving", 90),
    # fine
    (re.compile(r"Successfully exported|Export done|Model saved", re.I), "done", 100),
]

def _parse_line(line: str) -> tuple[str, int, str]:
    """Ritorna (phase, percent, message) dal testo di una riga stdout."""
    for pattern, phase, pct in _PATTERNS:
        if pattern.search(line):
            return phase, pct, line.strip()
    return "", -1, line.strip()

def _make_text_only_script(model_id: str, weight_format: str,
                            extra_flags: list, output_dir: str) -> str:
    """
    Genera uno script Python temporaneo per export text-only via API.
    Usa fn_get_submodels per escludere il vision encoder.
    Funziona per architetture senza _text variant nativa (LLaVA, InternVL, Phi3-V…).
    """
    extra = " ".join(extra_flags)
    return textwrap.dedent(f"""\
        import sys, logging
        logging.basicConfig(level=logging.INFO,
            format="%(asctime)s %(message)s", stream=sys.stdout)

        from optimum.exporters.openvino.__main__ import main_export

        # Cerca il sub-modello LLM nei pattern comuni per architettura VLM
        def text_only_submodels(model):
            for attr in ["language_model", "text_model", "text", "model"]:
                if hasattr(model, attr):
                    sub = getattr(model, attr)
                    # esclude il modello stesso e oggetti non-torch
                    import torch.nn as nn
                    if isinstance(sub, nn.Module) and sub is not model:
                        print(f"  ✂ escluso vision, uso: {{attr}}")
                        return {{"language_model": sub}}
            raise RuntimeError(
                f"Impossibile trovare il LLM in {{type(model).__name__}}. "
                "Architettura non supportata per potatura automatica."
            )

        main_export(
            model_name_or_path = {model_id!r},
            output             = {output_dir!r},
            task               = "text-generation-with-past",
            trust_remote_code  = True,
            fn_get_submodels   = text_only_submodels,
        )
        print("Export text-only completato.")
    """)


# ── Runner sincrono (da wrappare in QThread) ──────────────────────────────────
def run_conversion(
    job: ConversionJob,
    progress_cb: Callable[[ConversionProgress], None],
    cancelled_flag: Callable[[], bool]
) -> bool:
    """
    Esegue la conversione via subprocess.
    Se job.repo_id contiene '||text_only', smista verso la modalità
    text-only (CLI nativa o script Python con fn_get_submodels).
    """
    # ── parsing metadato text_only ────────────────────────────────────────────
    repo_id   = job.repo_id
    text_only = False
    if "||text_only" in repo_id:
        repo_id   = repo_id.replace("||text_only", "")
        text_only = True

    # ── risolvi quale CLI/Python usare ───────────────────────────────────────
    cli_path = _find_optimum_cli()
    if not cli_path:
        progress_cb(ConversionProgress(
            job_id=job.job_id, percent=0, phase="error",
            message="optimum-cli non trovato",
            error="Installa con: pip install optimum-intel[openvino]"
        ))
        return False

    job.output_dir.mkdir(parents=True, exist_ok=True)

    # ── text-only: scegli percorso CLI nativo vs script Python ───────────────
    if text_only:
        # leggi model_type dal config.json in cache HF (già scaricato)
        model_type = _get_model_type_from_args(job.cli_args)

        if _needs_python_script(model_type):
            # architettura senza _text variant → script con fn_get_submodels
            weight_fmt  = _get_arg_value(job.cli_args, "--weight-format") or "int4"
            extra_flags = _get_extra_flags(job.cli_args)
            script_src  = _make_text_only_script(
                repo_id, weight_fmt, extra_flags, str(job.output_dir)
            )
            return _run_python_script(
                script_src, job, repo_id, progress_cb, cancelled_flag
            )
        # altrimenti cade nel percorso CLI standard con task già sovrascritto

    # ── percorso CLI standard ─────────────────────────────────────────────────
    cli_args = list(job.cli_args)
    # rimpiazza repo_id nel caso contenesse il suffisso ||text_only
    for i, a in enumerate(cli_args):
        if "||text_only" in a:
            cli_args[i] = a.replace("||text_only", "")
    if cli_args and cli_args[0] == "optimum-cli":
        cli_args[0] = cli_path

    log.info(f"[{job.job_id}] avvio: {' '.join(cli_args[:6])}")
    progress_cb(ConversionProgress(
        job_id=job.job_id, percent=1, phase="loading",
        message=f"Avvio: {' '.join(cli_args[1:6])}..."
    ))
    return _run_subprocess(cli_args, job, progress_cb, cancelled_flag)


def _run_python_script(script_src: str, job: ConversionJob, repo_id: str,
                        progress_cb, cancelled_flag) -> bool:
    """Scrive uno script temporaneo e lo lancia con il Python dell'env."""
    import tempfile
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", prefix="ovforge_text_only_",
        delete=False, encoding="utf-8"
    ) as f:
        f.write(script_src)
        script_path = f.name

    log.info(f"[{job.job_id}] text-only via script: {script_path}")
    progress_cb(ConversionProgress(
        job_id=job.job_id, percent=1, phase="loading",
        message=f"✂ Solo LLM — export via fn_get_submodels per {repo_id}"
    ))

    py = _find_python()
    return _run_subprocess([py, script_path], job, progress_cb, cancelled_flag,
                            cleanup_file=script_path)


def _run_subprocess(cmd: list, job: ConversionJob,
                     progress_cb, cancelled_flag,
                     cleanup_file: str = None) -> bool:
    """Runner generico subprocess con progress parsing."""
    last_pct, last_phase = 0, "loading"
    try:
        proc = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            text=True, bufsize=1, encoding="utf-8", errors="replace"
        )
        stderr_lines = []
        for line in proc.stdout:
            if cancelled_flag():
                proc.terminate()
                progress_cb(ConversionProgress(
                    job_id=job.job_id, percent=last_pct, phase="error",
                    message="Cancellato", error="Cancelled"
                ))
                return False
            if not line.strip():
                continue
            stderr_lines.append(line)
            phase, pct, msg = _parse_line(line)
            if pct > last_pct: last_pct = pct
            if phase:           last_phase = phase
            progress_cb(ConversionProgress(
                job_id=job.job_id, percent=last_pct,
                phase=last_phase, message=msg
            ))
        proc.wait()

        if proc.returncode == 0:
            xml_files = list(job.output_dir.glob("*.xml"))
            if xml_files:
                progress_cb(ConversionProgress(
                    job_id=job.job_id, percent=100, phase="done",
                    message=f"✓ Completato → {job.output_dir}"
                ))
                return True
            err = "Nessun .xml trovato in output"
            progress_cb(ConversionProgress(
                job_id=job.job_id, percent=last_pct, phase="error",
                message=err, error=err
            ))
            return False
        else:
            tail = "".join(stderr_lines[-15:])
            progress_cb(ConversionProgress(
                job_id=job.job_id, percent=last_pct, phase="error",
                message=f"Errore (codice {proc.returncode})", error=tail
            ))
            return False
    except Exception as e:
        log.exception(f"[{job.job_id}] eccezione")
        progress_cb(ConversionProgress(
            job_id=job.job_id, percent=0, phase="error",
            message=str(e), error=str(e)
        ))
        return False
    finally:
        if cleanup_file:
            import os
            try: os.unlink(cleanup_file)
            except: pass


# ── Helper: cleanup output parziale ──────────────────────────────────────────
def _cleanup_partial_output(output_dir: Path) -> tuple[bool, int]:
    """
    Rimuove output_dir se contiene solo file parziali (0-byte o .xml senza .bin).
    Ritorna (rimosso: bool, n_file: int).
    Se la dir non esiste o contiene file integri, non tocca nulla.
    """
    if not output_dir.exists():
        return False, 0

    all_files = list(output_dir.rglob("*"))
    data_files = [f for f in all_files if f.is_file()]

    if not data_files:
        # dir vuota: rimuovi
        try:
            shutil.rmtree(output_dir)
            return True, 0
        except Exception as e:
            log.warning(f"cleanup: impossibile rimuovere dir vuota {output_dir}: {e}")
            return False, 0

    # xml con bin accoppiato → file integri
    xml_stems = {f.stem for f in data_files if f.suffix == ".xml" and f.stat().st_size > 0}
    bin_stems = {f.stem for f in data_files if f.suffix == ".bin" and f.stat().st_size > 0}
    paired_xml = xml_stems & bin_stems

    non_xml_intact = any(
        f.stat().st_size > 0 for f in data_files
        if f.suffix not in (".xml", ".bin")
    )

    if paired_xml or non_xml_intact:
        # output parzialmente integro: non rimuovere automaticamente
        log.info(f"cleanup: skip {output_dir} — contiene file integri")
        return False, 0

    # solo file parziali → rimuovi tutta la dir
    n = len(data_files)
    try:
        shutil.rmtree(output_dir)
        log.info(f"cleanup: rimossa {output_dir} ({n} file parziali)")
        return True, n
    except Exception as e:
        log.warning(f"cleanup: errore su {output_dir}: {e}")
        return False, 0


# ── Helper: estrai valori dagli args CLI ──────────────────────────────────────
def _get_arg_value(args: list, flag: str) -> Optional[str]:
    try:
        i = args.index(flag)
        return args[i + 1] if i + 1 < len(args) else None
    except ValueError:
        return None

def _get_model_type_from_args(args: list) -> str:
    """Legge model_type dal config.json della cache HF, dato il model_id negli args."""
    model_id = _get_arg_value(args, "--model") or ""
    if not model_id:
        return ""
    try:
        from huggingface_hub import hf_hub_download
        import json
        cfg_path = hf_hub_download(model_id, "config.json")
        with open(cfg_path) as f:
            return json.load(f).get("model_type", "")
    except Exception:
        return ""

def _get_extra_flags(args: list) -> list:
    """Estrae i flag extra (escludendo model, task, weight-format, output)."""
    skip_next = False
    skip_flags = {"--model", "--task", "--weight-format", "optimum-cli",
                  "export", "openvino"}
    result = []
    for i, a in enumerate(args):
        if skip_next:
            skip_next = False
            continue
        if a in skip_flags:
            skip_next = (a in {"--model", "--task", "--weight-format"})
            continue
        if not a.startswith("-") and i == len(args) - 1:
            continue   # ultimo argomento = output dir
        result.append(a)
    return result


# ── Downloader per modelli già convertiti (IR diretti) ────────────────────────
def run_download(
    repo_id: str,
    output_dir: Path,
    hf_token: Optional[str],
    progress_cb: Callable[[int, int, str], None],
    cancelled_flag: Callable[[], bool]
) -> bool:
    """
    Scarica un modello IR da HF hub con progress reale file per file.
    Supporta resume: i file già presenti e integri vengono saltati.
    progress_cb(downloaded_bytes, total_bytes, filename)
    """
    try:
        from huggingface_hub import HfApi, hf_hub_download
    except ImportError:
        progress_cb(0, 0, "ERRORE: huggingface_hub non installato")
        return False

    output_dir.mkdir(parents=True, exist_ok=True)
    api = HfApi(token=hf_token or None)

    # ── 1. Lista file nel repo con dimensioni ─────────────────────────────────
    try:
        progress_cb(0, 0, "Recupero lista file dal repo…")
        repo_files = list(api.list_repo_tree(repo_id, recursive=True))
    except Exception as e:
        log.exception(f"list_repo_tree fallita per {repo_id}")
        progress_cb(0, 0, f"ERRORE: impossibile listare il repo — {e}")
        return False

    if cancelled_flag():
        progress_cb(0, 0, "Annullato")
        return False

    # estrai solo i file (non le dir), con size lfs o size standard
    file_infos = []
    for item in repo_files:
        # RepoFile (huggingface_hub >= 0.21) ha sia .path che .rfilename
        path = getattr(item, "path", None) or getattr(item, "rfilename", None)
        if not path:
            continue
        lfs  = getattr(item, "lfs", None)
        size = lfs.size if lfs else (getattr(item, "size", 0) or 0)
        file_infos.append((path, int(size)))

    if not file_infos:
        progress_cb(0, 0, "ERRORE: nessun file trovato nel repo")
        return False

    total_bytes = sum(s for _, s in file_infos)
    downloaded_bytes = 0
    n = len(file_infos)

    # verifica una volta sola se hf_hub_download supporta local_dir_use_symlinks
    import inspect as _inspect
    _hf_dl_params = _inspect.signature(hf_hub_download).parameters
    _no_symlinks = {"local_dir_use_symlinks": False} if "local_dir_use_symlinks" in _hf_dl_params else {}

    log.info(f"Download {repo_id}: {n} file, {total_bytes/(1024**3):.2f} GB totali")

    # ── 2. Scarica file per file con resume ───────────────────────────────────
    for idx, (filename, file_size) in enumerate(file_infos):
        if cancelled_flag():
            progress_cb(downloaded_bytes, total_bytes, "Annullato")
            return False

        dest = output_dir / filename
        dest.parent.mkdir(parents=True, exist_ok=True)

        # resume: se il file esiste già con la dimensione corretta, salta
        if dest.exists() and file_size > 0 and dest.stat().st_size == file_size:
            log.info(f"  skip (già presente): {filename}")
            downloaded_bytes += file_size
            progress_cb(
                downloaded_bytes, total_bytes,
                f"[{idx+1}/{n}] ✓ skip {filename}"
            )
            continue

        # download con hf_hub_download → cache locale → copia in local_dir
        label = f"[{idx+1}/{n}] {filename}"
        progress_cb(downloaded_bytes, total_bytes, f"↓ {label}")
        try:
            cached = hf_hub_download(
                repo_id   = repo_id,
                filename  = filename,
                token     = hf_token or None,
                local_dir = str(output_dir),
                **_no_symlinks,   # file reali, non symlink — elimina = libera spazio
            )
        except Exception as e:
            log.warning(f"  download fallito {filename}: {e}")
            progress_cb(downloaded_bytes, total_bytes,
                        f"⚠ errore su {filename}: {e}")
            # non blocca — tenta gli altri file
            continue

        downloaded_bytes += file_size
        progress_cb(downloaded_bytes, total_bytes, f"✓ {label}")

    if cancelled_flag():
        progress_cb(downloaded_bytes, total_bytes, "Annullato")
        return False

    # ── 3. Verifica integrità finale ──────────────────────────────────────────
    missing = [
        fname for fname, fsize in file_infos
        if not (output_dir / fname).exists()
    ]
    if missing:
        msg = f"File mancanti dopo download: {', '.join(missing[:3])}"
        log.warning(msg)
        progress_cb(downloaded_bytes, total_bytes, f"⚠ {msg}")
        return False

    progress_cb(total_bytes, total_bytes,
                f"✓ Download completato ({n} file, {total_bytes/(1024**3):.2f} GB)")
    return True