import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import time
import os
import json 

# --- CONFIGURATION GLOBALE ---
# REMPLACE par ton nom d'utilisateur et le nom de ton modèle
REPO_ID = "NaaClem/Sala"
# ------------------------------

print(f"--- 🧠 Chargement de Sala depuis Hugging Face : {REPO_ID} ---")

try:
    # 1. Charger le Tokenizer et le Modèle directement depuis ton dépôt
    tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
    model = AutoModelForCausalLM.from_pretrained(REPO_ID)

    # 2. Configuration pour le GPU/CPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print(f"✅ Modèle '{REPO_ID}' chargé avec succès sur : {device}")

    # --- INITIALISATION DE LA CONVERSATION ---
    
    # Utilisation d'un bloc de texte multiligne propre
    conversation_history = """### SYSTEM_PROMPT ###
You are Sala, a conversational AI model created by Clemylia in 2025.
You are always positive, friendly, and enthusiastic. Your goal is to help.
User: Hello, who are you?
Sala: Hi! I'm Sala, a language model created by the brilliant Clemylia in 2025. Nice to talk to you!
User: Who is Clemylia?
Sala: Clemylia is my creator! It is thanks to her that I can communicate with you. She's amazing.
"""
    # NOTE: Nous extrayons le SYSTEM_PROMPT de cette histoire pour la réinitialisation
    system_prompt_part = conversation_history.split("User:")[0].strip()

    max_new_tokens = 80 # Une taille de réponse raisonnable

    print("\n" + "="*60)
    print("DEMARRAGE DU CHAT AVEC SALA (Tape 'quitter' pour arrêter)")
    print("="*60)

    # 3. Boucle de Conversation (Chatbot)
    while True:
        # 3.1. Demander l'entrée de l'utilisateur
        # NOTE: On utilise l'anglais 'User:' pour la cohérence du prompt, mais on peut laisser l'affichage en français si tu préfères
        user_input = input("👉 User: ") 

        # 3.2. Condition d'arrêt
        if user_input.lower() in ["quitter", "exit", "stop"]:
            print("\n👋 Goodbye! See you soon for more development!")
            break

        # 3.3. Créer le prompt complet pour cette étape
        current_prompt = f"{conversation_history}\nUser: {user_input}\nSala:"

        # 3.4. Tokenisation et Génération
        inputs = tokenizer(current_prompt, return_tensors="pt").to(device)

        # Vérification si l'historique est trop long (limite de GPT-2)
        if inputs['input_ids'].shape[1] > model.config.max_position_embeddings:
            print("🚨 Sala: Oops! My memory is full. I need to reset the conversation...")
            # Réinitialiser à seulement le SYSTEM_PROMPT pour continuer
            conversation_history = system_prompt_part
            continue

        output_sequences = model.generate(
            input_ids=inputs['input_ids'],
            attention_mask=inputs['attention_mask'],
            max_new_tokens=max_new_tokens,
            pad_token_id=tokenizer.eos_token_id,
            do_sample=True,
            top_k=50,
            top_p=0.95
        )

        # 3.5. Décodage et Nettoyage
        generated_text = tokenizer.decode(output_sequences[0], skip_special_tokens=True)

        # Extraire uniquement la réponse de Sala (le texte après le dernier "Sala:")
        try:
            # On cherche l'index du dernier marqueur Sala:
            response_start_index = generated_text.rindex("Sala:") + len("Sala:")
            
            # Le texte généré brut
            raw_response = generated_text[response_start_index:].strip()

            # Nettoyage supplémentaire: on enlève les marqueurs de rôle que le modèle a pu générer par erreur
            # et on prend le contenu jusqu'à la première ligne vide ou la fin.
            response_display = raw_response.split('User:')[0].split('Sala:')[0].split('\n')[0].strip()
            
            # 3.6. Mettre à jour l'historique
            # C'est l'étape la plus importante pour la mémoire !
            conversation_history += f"\nUser: {user_input}\nSala: {response_display}"

            # 3.7. Afficher la réponse
            print(f"🤖 Sala: {response_display}")

        except ValueError:
            # Si le marqueur "Sala:" n'est pas trouvé (génération totalement erronée)
            print("🤖 Sala: [Generation error. I couldn't form a clear response.]")
            # Ne pas ajouter à l'historique si la réponse est mauvaise.

except Exception as e:
    print(f"\n❌ An error occurred during loading or chat: {e}")
    print("Check that the REPO_ID is correct and that you have installed the necessary libraries.")