# ============================================
# 1. Imports
# ============================================
import torch
import gradio as gr
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
)
from peft import PeftModel
from huggingface_hub import login
import os

# ============================================
# 2. Configurations principales
# ============================================
# Les noms de modèles doivent être des chemins relatifs ou des identifiants de Hugging Face Hub
base_model = "google/gemma-2-2b-it"
# Le répertoire où le modèle finetuné sera sauvegardé/chargé
# Dans un Gradio Space, on le place souvent à la racine
fine_tuned_model_dir = "gemma-brvm-lora-v2"

# Assure-toi que le modèle finetuné existe sur le Hub ou dans le Space
# Sinon, cette partie échouera.
model_id = "votre-utilisateur/votre-nouveau-modele-lora" # Remplace par l'ID de ton modèle finetuné
# Pour un déploiement, on ne re-entraîne pas, on charge le modèle déjà entraîné.
# Le code d'entraînement est généralement dans un script séparé ou sur une autre plateforme.

# ============================================
# 3. Connexion à Hugging Face
# ============================================
# Pour un Gradio Space, la connexion est gérée par les secrets.
# Si le modèle est privé, assure-toi que la variable HF_TOKEN est définie.
try:
    login()
except Exception:
    print("Connexion Hugging Face impossible. Le script continuera si les ressources sont publiques.")

# ============================================
# 4. Fonction de chargement du modèle
# ============================================
def load_model_and_tokenizer():
    """
    Charge le modèle de base avec la configuration de quantization et applique l'adaptateur LoRA.
    """
    print("⏳ Chargement du modèle...")

    # Configuration 4-bit (QLoRA)
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16,
        bnb_4bit_use_double_quant=True,
    )

    # Charger le modèle de base avec la quantization
    base_model_loaded = AutoModelForCausalLM.from_pretrained(
        base_model,
        quantization_config=bnb_config,
        device_map="auto",
        token=True
    )

    # Charger le tokenizer
    tokenizer = AutoTokenizer.from_pretrained(base_model, token=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    # Charger l'adaptateur LoRA
    # C'est la partie cruciale pour le déploiement.
    model = PeftModel.from_pretrained(base_model_loaded, model_id)

    print("✅ Modèle et tokenizer chargés et adaptés.")
    return model, tokenizer

# Charger le modèle globalement pour qu'il ne soit chargé qu'une seule fois
# au démarrage du Space.
model, tokenizer = load_model_and_tokenizer()

# ============================================
# 5. Fonction de génération de texte
# ============================================
def generate_response(prompt_text, max_length=512, temperature=0.7):
    """
    Génère une réponse à partir d'un prompt utilisateur.
    """
    inputs = tokenizer(
        f"Instruction: {prompt_text}\nRéponse:",
        return_tensors="pt"
    ).to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_length,
            do_sample=True,
            temperature=temperature,
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Nettoyage de la réponse pour ne retourner que la partie générée
    # La ligne suivante suppose que le modèle répond après "Réponse:"
    response = response.split("Réponse:")[-1].strip()

    return response

# ============================================
# 6. Déploiement de l'interface Gradio
# ============================================
# Création de l'interface Gradio
title = "Modèle Gemma Fine-tuné sur le BRVM"
description = "Entrez une question sur la Bourse Régionale des Valeurs Mobilières (BRVM) et obtenez une réponse."

iface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(lines=5, label="Votre question"),
        gr.Slider(minimum=1, maximum=1024, step=1, value=512, label="Longueur maximale de la réponse"),
        gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.7, label="Température (créativité)"),
    ],
    outputs="text",
    title=title,
    description=description,
    live=True, # Optionnel : mise à jour en direct
    theme=gr.themes.Monochrome() # Un thème sympa
)

# Lancement de l'interface
# Le `share=False` est important pour un Space.
iface.launch(share=False)