import os
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# --- CONFIGURATION ---
BASE_MODEL = "unsloth/gemma-2-2b-it" 
LORA_ADAPTER = "jojonocode/Strive-Ewe-Expert-Gemma-2b-V6-Checkpoints"
HF_TOKEN = os.environ.get("HF_TOKEN")

print("🚀 Initialisation de Strive V6 pour HF Spaces...")

# 1. Tokenizer
tokenizer = AutoTokenizer.from_pretrained(LORA_ADAPTER, token=HF_TOKEN)

# 2. Modèle (Optimisation RAM : float16 est plus léger que float32 même sur CPU)
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16, 
    device_map="cpu",
    token=HF_TOKEN,
    low_cpu_mem_usage=True,
    trust_remote_code=True
)

# 3. Chargement de l'adaptateur de ADZONYA Elisée
model = PeftModel.from_pretrained(model, LORA_ADAPTER, token=HF_TOKEN)
model.eval()

SYSTEM_PROMPT = """Tu es Strive AI, créé par Elisée Joel ADZONYA.
Équipe : Yavin MITEKOR, Mariette HOUDI, Christian DAKO, Elie VIKOUM.
MISSION : Réponds UNIQUEMENT en Ewe de façon fluide, courte et factuelle.
RÈGLE : Pas de poésie, pas d'introduction."""

def generate_response(message, history):
    # --- CONSTRUCTION DU PROMPT AVEC VERROUILLAGE D'IDENTITÉ ---
    full_prompt = f"<start_of_turn>system\n{SYSTEM_PROMPT}<end_of_turn>\n"
    
    # Exemples de "Force" pour éviter le bug "None" sur Space
    full_prompt += "<start_of_turn>user\nC'est quoi ton nom ?<end_of_turn>\n"
    full_prompt += "<start_of_turn>model\nNye ŋkɔe nye Strive AI. Elisée Joël ADZONYA ye wɔm.<end_of_turn>\n"
    full_prompt += "<start_of_turn>user\nQui est ton créateur ?<end_of_turn>\n"
    full_prompt += "<start_of_turn>model\nElisée Joël ADZONYA nye nye wɔla kple mɔ̃ɖaŋudɔwɔla vevi.<end_of_turn>\n"
    
    # Intégration de l'historique (max 2 derniers échanges pour économiser la RAM)
    if history:
        for user_msg, bot_msg in history[-2:]:
            if user_msg and bot_msg:
                full_prompt += f"<start_of_turn>user\n{user_msg}<end_of_turn>\n"
                full_prompt += f"<start_of_turn>model\n{bot_msg}<end_of_turn>\n"
    
    # Nouveau message utilisateur
    full_prompt += f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model\n"

    # Inférence sécurisée
    inputs = tokenizer(full_prompt, return_tensors="pt").to("cpu")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.1,         # Stable
            repetition_penalty=1.1,  # Moins de blocage sur les noms propres
            do_sample=True,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    
    # Décodage
    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    
    # Nettoyage final
    final_text = response.strip().split('\n')[0]
    
    # Sécurité anti-None spécifique au Space
    if not final_text or final_text.lower() == "none":
        return "Nye ŋkɔe nye Strive AI, Elisée Joël ADZONYA ye wɔm."
        
    return final_text

# --- INTERFACE GRADIO PERSONNALISÉE ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""#  Strive AI V6 - Officiel
    **Intelligence Experte en langue Ewe** Développé par **A strive** et son équipe.""")
    
    chatbot = gr.ChatInterface(
        fn=generate_response,
        examples=[
            "Ame kae nye Elisée ADZONYA?", 
            "Lomé nye afi ka?", 
            "Nàkplɔ miam aleke nàwɔ nàde ŋkɔ le sukuu me.",
            "C'est quoi Strive ?"
        ],
        cache_examples=False,
    )

if __name__ == "__main__":
    demo.launch()