#!/usr/bin/env python
"""Verifie l'environnement InkFlow et pre-telecharge les modeles MLX.

Usage :
    python scripts/setup_models.py            # tout verifier + telecharger
    python scripts/setup_models.py --check    # verifier sans telecharger

Pre-requis systeme : Apple Silicon, Python >= 3.11, ffmpeg (brew install ffmpeg).
"""
from __future__ import annotations

import argparse
import platform
import shutil
import sys

# Permet de lancer le script directement depuis backend/.
sys.path.insert(0, str(__import__("pathlib").Path(__file__).resolve().parents[1]))

from inkflow.config import (  # noqa: E402
    GEMMA_MODEL,
    KOKORO_MODEL,
    QWEN3_TTS_MODEL,
    ensure_dirs,
)


def check_env() -> bool:
    ok = True
    print(f"• Plateforme       : {platform.platform()} ({platform.machine()})")
    if platform.machine() != "arm64":
        print("  ! Attendu arm64 (Apple Silicon) — MLX ne sera pas optimal.")
    print(f"• Python           : {sys.version.split()[0]}")
    if sys.version_info < (3, 11):
        print("  ! Python >= 3.11 requis."); ok = False

    for mod in ("mlx", "mlx_lm", "mlx_audio", "ebooklib", "bs4",
                "soundfile", "mutagen", "fastapi"):
        try:
            __import__(mod)
            print(f"• import {mod:12s}: OK")
        except Exception as exc:  # noqa: BLE001
            print(f"• import {mod:12s}: ECHEC ({exc})"); ok = False

    ff = shutil.which("ffmpeg")
    print(f"• ffmpeg           : {ff or 'INTROUVABLE — brew install ffmpeg'}")
    ok = ok and bool(ff)
    return ok


def download_lm(model_id: str) -> None:
    from mlx_lm import load
    print(f"  -> LM   {model_id}")
    load(model_id)


def download_tts(model_id: str) -> None:
    from mlx_audio.tts.utils import load_model
    print(f"  -> TTS  {model_id}")
    load_model(model_id)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--check", action="store_true", help="verifier sans telecharger")
    args = ap.parse_args()

    ensure_dirs()
    print("== Verification de l'environnement ==")
    env_ok = check_env()

    if args.check:
        return 0 if env_ok else 1
    if not env_ok:
        print("\nEnvironnement incomplet — corrige les points ci-dessus avant de continuer.")
        return 1

    print("\n== Telechargement des modeles (peut etre long la 1re fois) ==")
    download_lm(GEMMA_MODEL)
    download_tts(KOKORO_MODEL)
    download_tts(QWEN3_TTS_MODEL)
    print("\nTout est pret.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())