app.py

# app.py
from __future__ import annotations

import json
import os
import threading
import time
from dataclasses import dataclass, asdict
from typing import Any, Dict, Generator, List, Optional

from flask import Flask, Response, jsonify, request, stream_with_context
from flask_cors import CORS
from dotenv import load_dotenv

load_dotenv()  # charge .env si présent


# -------------------------------
# Dépendances externes
# -------------------------------
try:
    from llama_cpp import Llama
except Exception as e:  # pragma: no cover
    raise RuntimeError(
        "llama-cpp-python is required. Install with `pip install --upgrade llama-cpp-python`."
    ) from e

try:
    from huggingface_hub import hf_hub_download
    HAS_HF = True
except Exception:
    HAS_HF = False


# -------------------------------
# Helpers ENV
# -------------------------------
def getenv_str(name: str, default: Optional[str] = None) -> Optional[str]:
    v = os.getenv(name, default)
    return v if (v is None or v.strip() != "") else default


def getenv_int(name: str, default: int) -> int:
    try:
        return int(os.getenv(name, default))
    except Exception:
        return default


def getenv_bool(name: str, default: bool = False) -> bool:
    v = os.getenv(name)
    if v is None:
        return default
    return v.strip().lower() in ("1", "true", "yes", "y", "on")


# -------------------------------
# Modèle & Configuration
# -------------------------------
@dataclass
class ModelConfig:
    # Chemin local vers un GGUF (sous MODELS_DIR)
    model_path: Optional[str] = getenv_str("MODEL_PATH")

    # Paramètres d'inférence
    n_ctx: int = getenv_int("N_CTX", 4096)
    n_threads: int = getenv_int("N_THREADS", os.cpu_count() or 4)
    n_batch: int = getenv_int("N_BATCH", 256)
    n_gpu_layers: int = getenv_int("N_GPU_LAYERS", 0)  # CPU par défaut
    verbose: bool = os.getenv("LLM_VERBOSE", "0") == "1"

    def to_display(self) -> Dict[str, Any]:
        d = asdict(self)
        if d.get("model_path"):
            d["model_path"] = os.path.basename(str(d["model_path"]))
        return d


class ModelManager:
    """Gestion thread-safe du modèle (lazy load + reload)."""

    def __init__(self, config: Optional[ModelConfig] = None):
        self._lock = threading.RLock()
        self._llm: Optional[Llama] = None
        self._config = config or ModelConfig()

    @property
    def config(self) -> ModelConfig:
        return self._config

    def is_loaded(self) -> bool:
        with self._lock:
            return self._llm is not None

    def set_config(self, new_cfg: ModelConfig) -> None:
        """Met à jour la config sans charger en mémoire."""
        with self._lock:
            self._config = new_cfg
            # ne touche pas à _llm

    def get_llm(self) -> Llama:
        with self._lock:
            if self._llm is not None:
                return self._llm
            self._llm = self._load_model(self._config)
            return self._llm

    def _load_model(self, cfg: ModelConfig) -> Llama:
        if not cfg.model_path or not os.path.exists(cfg.model_path):
            raise RuntimeError(
                f"MODEL_PATH does not exist: {cfg.model_path!r}. Use /load to fetch a model or set MODEL_PATH."
            )
        t0 = time.time()
        llm = Llama(
            model_path=cfg.model_path,
            n_ctx=cfg.n_ctx,
            n_threads=cfg.n_threads,
            n_batch=cfg.n_batch,
            n_gpu_layers=cfg.n_gpu_layers,
            verbose=cfg.verbose,
        )
        dt = time.time() - t0
        print(f"[LLM] Loaded model in {dt:.2f}s | cfg={cfg.to_display()}")
        return llm

    def reload(self, new_cfg: Optional[ModelConfig] = None) -> Dict[str, Any]:
        with self._lock:
            self._llm = None
            if new_cfg is not None:
                self._config = new_cfg
            _ = self.get_llm()  # force le chargement pour valider
            return {"status": "reloaded", "config": self._config.to_display()}


# -------------------------------
# Flask app
# -------------------------------
app = Flask(__name__)
CORS(app)

MODELS_DIR = getenv_str("MODELS_DIR", "/models") or "/models"
os.makedirs(MODELS_DIR, exist_ok=True)

# Options d'auto-load au démarrage
AUTO_REPO_ID = getenv_str("AUTO_REPO_ID")
AUTO_FILENAME = getenv_str("AUTO_FILENAME")  # ex: "gemma-3-270m-it-Q4_K_M.gguf"
AUTO_IF_MISSING = getenv_bool("AUTO_IF_MISSING", True)  # ne télécharge que si le fichier manque
AUTO_PRELOAD = getenv_bool("AUTO_PRELOAD", False)  # charge en mémoire au boot

manager = ModelManager()


def _bootstrap_autoload() -> Dict[str, Any]:
    """
    Si MODEL_PATH n'existe pas et que AUTO_REPO_ID/AUTO_FILENAME sont fournis,
    tente un téléchargement dans MODELS_DIR, puis prépare MODEL_PATH.
    Si AUTO_PRELOAD=True, charge aussi le modèle en mémoire (via reload()).
    """
    info: Dict[str, Any] = {
        "models_dir": MODELS_DIR,
        "auto_repo_id": AUTO_REPO_ID,
        "auto_filename": AUTO_FILENAME,
        "auto_if_missing": AUTO_IF_MISSING,
        "auto_preload": AUTO_PRELOAD,
        "performed": False,
        "path": None,
        "error": None,
    }

    cfg = manager.config
    target_path = cfg.model_path

    # Si MODEL_PATH n'est pas défini mais qu'on a un AUTO_FILENAME, on propose une cible.
    if (not target_path) and AUTO_FILENAME:
        target_path = os.path.join(MODELS_DIR, os.path.basename(AUTO_FILENAME))

    # Si on a déjà un fichier présent, pas besoin de télécharger.
    if target_path and os.path.exists(target_path):
        # Met à jour la config (sans précharger)
        manager.set_config(
            ModelConfig(
                model_path=target_path,
                n_ctx=cfg.n_ctx,
                n_threads=cfg.n_threads,
                n_batch=cfg.n_batch,
                n_gpu_layers=cfg.n_gpu_layers,
                verbose=cfg.verbose,
            )
        )
        if AUTO_PRELOAD:
            try:
                manager.reload()  # charge en mémoire
            except Exception as e:
                info["error"] = f"preload_failed: {e}"
        info["performed"] = False
        info["path"] = target_path
        print(f"[BOOT] Model file already present at {target_path}")
        return info

    # Ici, le fichier n'existe pas : on tente un auto-load si possible.
    if AUTO_IF_MISSING and AUTO_REPO_ID and AUTO_FILENAME and HAS_HF:
        try:
            print(f"[BOOT] Auto-load from HF: {AUTO_REPO_ID}/{AUTO_FILENAME}")
            local_cached = hf_hub_download(repo_id=AUTO_REPO_ID, filename=AUTO_FILENAME)
            target_path = os.path.join(MODELS_DIR, os.path.basename(AUTO_FILENAME))
            if os.path.abspath(local_cached) != os.path.abspath(target_path):
                # Copie chunkée
                with open(local_cached, "rb") as src, open(target_path, "wb") as dst:
                    while True:
                        chunk = src.read(1024 * 1024)
                        if not chunk:
                            break
                        dst.write(chunk)

            new_cfg = ModelConfig(
                model_path=target_path,
                n_ctx=cfg.n_ctx,
                n_threads=cfg.n_threads,
                n_batch=cfg.n_batch,
                n_gpu_layers=cfg.n_gpu_layers,
                verbose=cfg.verbose,
            )

            if AUTO_PRELOAD:
                manager.reload(new_cfg)  # charge en mémoire tout de suite
            else:
                manager.set_config(new_cfg)  # défini le chemin, chargera à la 1re requête

            info["performed"] = True
            info["path"] = target_path
            print(f"[BOOT] Auto-loaded model to {target_path} (preload={AUTO_PRELOAD})")
            return info
        except Exception as e:
            info["error"] = str(e)
            print(f"[BOOT][ERROR] Auto-load failed: {e}")
            return info

    # Pas d'auto-load possible / activé
    return info


# Effectue l'auto-load au démarrage (non bloquant au niveau de la disponibilité HTTP)
BOOT_INFO = _bootstrap_autoload()


# -------------------------------
# Endpoints d'info
# -------------------------------
@app.get("/health")
def health() -> Any:
    """
    Health check "soft":
    - Si MODEL_PATH absent ou fichier manquant -> status="unloaded" (HTTP 200), sans 500.
    - Si fichier présent mais modèle non encore chargé -> status="available".
    - Si modèle en mémoire -> status="ready".
    """
    try:
        cfg = manager.config
        path = cfg.model_path
        if not path or not os.path.exists(path):
            return jsonify(
                {
                    "status": "unloaded",
                    "reason": "MODEL_PATH is not set or file not found",
                    "model": cfg.to_display(),
                    "exists": bool(path and os.path.exists(path)),
                    "boot_info": BOOT_INFO,
                    "hint": 'Use /load then /reload with {"model_path": "/models/your.gguf"} '
                            'or set AUTO_REPO_ID/AUTO_FILENAME in .env for auto-load.',
                }
            )
        if not manager.is_loaded():
            # Fichier OK, modèle pas encore chargé
            return jsonify(
                {
                    "status": "available",
                    "model": cfg.to_display(),
                    "path": path,
                    "boot_info": BOOT_INFO,
                    "hint": 'Call /reload with {"model_path": "%s"}' % path,
                }
            )
        # Modèle déjà en mémoire
        return jsonify({"status": "ready", "model": cfg.to_display(), "path": path})
    except Exception as e:
        # Cas vraiment inattendu
        return jsonify({"status": "error", "error": str(e)}), 500


@app.get("/models")
def models() -> Any:
    """
    Liste les GGUF disponibles et l'état du modèle actif.
    """
    cfg = manager.config

    available = []
    for f in os.listdir(MODELS_DIR):
        if f.lower().endswith(".gguf"):
            full = os.path.join(MODELS_DIR, f)
            try:
                size = os.path.getsize(full)
            except Exception:
                size = None
            available.append({"name": f, "path": full, "size_bytes": size})

    active_path = cfg.model_path
    active_exists = bool(active_path and os.path.exists(active_path))
    state = (
        "ready"
        if manager.is_loaded()
        else ("available" if active_exists else "unloaded")
    )

    return jsonify(
        {
            "object": "list",
            "active": {
                "path": active_path,
                "basename": os.path.basename(active_path) if active_path else None,
                "exists": active_exists,
                "loaded": manager.is_loaded(),
                "state": state,
                "config": cfg.to_display(),
            },
            "available": available,
            "hint": 'Switch with POST /reload {"model_path": "/models/<file>.gguf"}',
            "boot_info": BOOT_INFO,
        }
    )


# -------------------------------
# Gestion des modèles (load/reload)
# -------------------------------
@app.post("/load")
def load_model_file() -> Any:
    """
    Télécharge un GGUF depuis Hugging Face Hub vers MODELS_DIR.
    Body: {"repo_id": str, "filename": str}
    Ne recharge pas automatiquement le modèle — appeler /reload ensuite.
    """
    if not HAS_HF:
        return (
            jsonify(
                {"error": {"message": "huggingface_hub is not installed on server."}}
            ),
            500,
        )

    body = request.get_json(force=True)
    repo_id = (body.get("repo_id") or "").strip()
    filename = (body.get("filename") or "").strip()
    if not repo_id or not filename:
        return jsonify({"error": {"message": "repo_id and filename are required"}}), 400

    try:
        # Respecte HF_TOKEN éventuellement défini dans l'env.
        local_cached = hf_hub_download(repo_id=repo_id, filename=filename)
        target_path = os.path.join(MODELS_DIR, os.path.basename(filename))
        if os.path.abspath(local_cached) != os.path.abspath(target_path):
            # Copie (chunkée) depuis le cache HF vers MODELS_DIR
            with open(local_cached, "rb") as src, open(target_path, "wb") as dst:
                while True:
                    chunk = src.read(1024 * 1024)
                    if not chunk:
                        break
                    dst.write(chunk)
        return jsonify(
            {
                "status": "downloaded",
                "repo_id": repo_id,
                "filename": filename,
                "path": target_path,
                "hint": 'POST /reload {"model_path": "%s"}' % target_path,
            }
        )
    except Exception as e:
        return jsonify({"error": {"message": str(e)}}), 500


@app.post("/reload")
def reload_model() -> Any:
    """
    Recharge le modèle avec un nouveau MODEL_PATH et/ou nouveaux hyperparamètres.
    Body: {"model_path": str, "n_ctx": int, "n_threads": int, "n_batch": int, "n_gpu_layers": int, "verbose": bool}
    """
    payload = request.get_json(silent=True) or {}
    model_path = payload.get("model_path", manager.config.model_path)

    cfg = ModelConfig(
        model_path=model_path,
        n_ctx=int(payload.get("n_ctx", manager.config.n_ctx)),
        n_threads=int(payload.get("n_threads", manager.config.n_threads)),
        n_batch=int(payload.get("n_batch", manager.config.n_batch)),
        n_gpu_layers=int(payload.get("n_gpu_layers", manager.config.n_gpu_layers)),
        verbose=bool(payload.get("verbose", manager.config.verbose)),
    )
    try:
        info = manager.reload(cfg)
        return jsonify(info)
    except Exception as e:
        return jsonify({"status": "error", "error": str(e)}), 400


# -------------------------------
# Chat Completions (OpenAI-like) + SSE
# -------------------------------
def sse_format(data: Dict[str, Any]) -> str:
    return "data: " + json.dumps(data, ensure_ascii=False) + "\n\n"


@app.post("/v1/chat/completions")
def chat_completions() -> Any:
    """
    Endpoint OpenAI-like (non-stream ou SSE si {"stream": true}).
    Body minimal:
      {
        "messages": [{"role":"system","content":"..."}, {"role":"user","content":"..."}],
        "stream": true|false,
        "temperature": 0.7,
        "top_p": 0.95,
        "top_k": 64,
        "min_p": 0.001,
        "max_tokens": 256,
        "stop": ["..."],
        "response_format": {...}
      }
    """
    llm = manager.get_llm()
    body = request.get_json(force=True)

    messages: List[Dict[str, str]] = body.get("messages", [])
    if not messages:
        return jsonify({"error": {"message": "messages[] is required"}}), 400

    stream = bool(body.get("stream", False))

    params: Dict[str, Any] = {}
    for key in ("temperature", "top_p", "top_k", "min_p", "max_tokens", "stop"):
        if key in body:
            params[key] = body[key]

    if "response_format" in body:
        params["response_format"] = body["response_format"]

    if not stream:
        t0 = time.time()
        resp = llm.create_chat_completion(messages=messages, **params)
        dt = time.time() - t0
        out = {
            "id": resp.get("id", f"chatcmpl-{int(time.time()*1000)}"),
            "object": "chat.completion",
            "created": int(time.time()),
            "model": os.path.basename(manager.config.model_path or "local.gguf"),
            "choices": resp.get("choices", []),
            "usage": resp.get("usage", {}),
            "latency_ms": int(dt * 1000),
        }
        return jsonify(out)

    # Streaming SSE
    def generate() -> Generator[str, None, None]:
        t0 = time.time()
        try:
            for chunk in llm.create_chat_completion(
                messages=messages, stream=True, **params
            ):
                # Chunks similaires au format OpenAI
                out = {
                    "id": chunk.get("id", f"chatcmpl-{int(time.time()*1000)}"),
                    "object": "chat.completion.chunk",
                    "created": int(time.time()),
                    "model": os.path.basename(manager.config.model_path or "local.gguf"),
                    "choices": chunk.get("choices", []),
                }
                yield sse_format(out)
        except Exception as e:
            yield sse_format({"error": {"message": str(e)}})
        finally:
            yield sse_format({"done": True})
            yield "data: [DONE]\n\n"

    return Response(stream_with_context(generate()), mimetype="text/event-stream")


# -------------------------------
# Extraction JSON (non-stream)
# -------------------------------
@app.post("/v1/extract")
def extract() -> Any:
    """
    Extraction JSON avec schéma optionnel.
    Body:
      - text: str (requis)
      - schema: dict (optionnel) -> transmis via response_format si supporté
      - system_prompt: str (optionnel)
      - temperature/top_p/top_k/min_p/max_tokens/stop: optionnels
    """
    llm = manager.get_llm()
    body = request.get_json(force=True)

    text: str = body.get("text", "").strip()
    if not text:
        return jsonify({"error": {"message": "'text' is required"}}), 400

    schema = body.get("schema")
    system_prompt = body.get(
        "system_prompt",
        "You are a helpful assistant that outputs ONLY valid minified JSON. Do not include prose.",
    )

    params: Dict[str, Any] = {}
    for key in ("temperature", "top_p", "top_k", "min_p", "max_tokens", "stop"):
        if key in body:
            params[key] = body[key]

    if isinstance(schema, dict):
        params["response_format"] = {"type": "json_object", "schema": schema}

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": text},
    ]

    try:
        t0 = time.time()
        resp = llm.create_chat_completion(messages=messages, **params)
        dt = time.time() - t0
        content = ((resp.get("choices") or [{}])[0].get("message") or {}).get(
            "content", ""
        )
        try:
            parsed = json.loads(content)
        except Exception as pe:
            return (
                jsonify(
                    {
                        "error": {
                            "message": "Model did not produce valid JSON.",
                            "model_output": content,
                            "parse_error": str(pe),
                        }
                    }
                ),
                422,
            )

        # Contrôle léger des clés attendues
        if isinstance(schema, dict) and "properties" in schema:
            missing = [
                k
                for k in schema["properties"].keys()
                if not (isinstance(parsed, dict) and k in parsed)
            ]
            if missing:
                return (
                    jsonify(
                        {
                            "error": {
                                "message": "JSON is valid but missing required keys (best-effort check).",
                                "missing": missing,
                                "model_output": parsed,
                            }
                        }
                    ),
                    206,
                )

        return jsonify(
            {
                "object": "extract.result",
                "created": int(time.time()),
                "model": os.path.basename(manager.config.model_path or "local.gguf"),
                "latency_ms": int(dt * 1000),
                "data": parsed,
            }
        )
    except Exception as e:
        return jsonify({"error": {"message": str(e)}}), 500


# -------------------------------
# Entrypoint
# -------------------------------
if __name__ == "__main__":
    host = getenv_str("HOST", "0.0.0.0") or "0.0.0.0"
    port = getenv_int("PORT", 8000)
    debug = os.getenv("FLASK_DEBUG", "0") == "1"
    print("Starting Flask on", host, port, "| MODELS_DIR=", MODELS_DIR)
    app.run(host=host, port=port, debug=debug)