import os
import sys
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Initialisation de l'application FastAPI
app = FastAPI(
    title="Makoa Learning AI - Engine Q3 Stable",
    description="Moteur d'inférence optimisé basé sur Gemma 3 4B IT (TensorBlock GGUF)"
)

# -------------------------------------------------------------------------
# CONFIGURATION DU MODÈLE : Version standard TensorBlock (Sans bug QAT)
# -------------------------------------------------------------------------
MODEL_REPO = "tensorblock/gemma-3-4b-it-GGUF"
MODEL_FILE = "gemma-3-4b-it-Q3_K_M.gguf"  # Version 3-bits ultra-rapide sur CPU

# -------------------------------------------------------------------------
# ÉTAPE 1 : TÉLÉCHARGEMENT SÉCURISÉ DEPUIS HUGGING FACE
# -------------------------------------------------------------------------
print("=== [ÉTAPES 1/3] Récupération du modèle stable depuis TensorBlock ===", flush=True)
try:
    hf_token = os.environ.get("HF_TOKEN")
    if not hf_token:
        raise ValueError(
            "Le secret 'HF_TOKEN' est introuvable dans les paramètres de ton Space. "
            "Vérifie l'onglet Settings."
        )

    # Téléchargement de la version stable Q3_K_M (Environ 2.0 Go)
    model_path = hf_hub_download(
        repo_id=MODEL_REPO, 
        filename=MODEL_FILE,
        token=hf_token
    )
    print(f"✅ Modèle récupéré avec succès et localisé ici : {model_path}", flush=True)
except Exception as e:
    print(f"❌ Erreur critique lors de la récupération du modèle : {e}", flush=True)
    raise e

# -------------------------------------------------------------------------
# ÉTAPE 2 : INITIALISATION DU MOTEUR LLAMA.CPP
# -------------------------------------------------------------------------
print("=== [ÉTAPES 2/3] Chargement du modèle en RAM ===", flush=True)
try:
    llm = Llama(
        model_path=model_path,
        n_ctx=2048,      # Taille de contexte idéale pour les 16 Go de RAM gratuits
        n_threads=2,     # Utilisation stricte des 2 cœurs vCPU disponibles
        verbose=False    # Nettoie les logs internes pour éviter les écritures inutiles
    )
    print("=== [ÉTAPES 3/3] Backend opérationnel et prêt à l'écoute ! ===", flush=True)
except Exception as e:
    print(f"❌ Échec de l'initialisation du moteur llama.cpp : {e}", flush=True)
    raise e

# -------------------------------------------------------------------------
# SCHÉMAS DE DONNÉES ET ROUTES API
# -------------------------------------------------------------------------
class ChatRequest(BaseModel):
    prompt: str
    max_tokens: int = 256
    temperature: float = 0.3
    top_p: float = 0.9

@app.get("/")
def health_check():
    """Route technique indispensable pour valider le statut du Space."""
    return {
        "status": "online", 
        "model_loaded": MODEL_FILE,
        "mode": "tensorblock_standard_q3"
    }

@app.post("/v1/chat")
def ask_gemma(request: ChatRequest):
    """Point d'accès pour générer tes missions, dialogues et traductions."""
    print(f"📥 Requête reçue ! Début du prompt : {request.prompt[:100]}...", flush=True)
    try:
        # Inférence directe sur le CPU du Space
        response = llm(
            prompt=request.prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            stop=["<end_of_turn>", "<eos>", "\nUser:", "\n\nUser:"]
        )
        
        return {
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": response["choices"][0]["text"]
                    }
                }
            ],
            "usage": response.get("usage", {})
        }
    except Exception as e:
        print(f"❌ Erreur pendant l'inférence : {str(e)}", flush=True)
        raise HTTPException(
            status_code=500, 
            detail=f"Erreur d'inférence : {str(e)}"
        )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=False)