# Prérequis :
# Installation de llala.cpp :
#  pip install --no-cache-dir llama-cpp-python
#
# Téléchargement des modèles :
#  mkdir Models
#  cd Models
#
# Première possibilité :
#  huggingface-cli download bartowski/Qwen2.5-1.5B-Instruct-GGUF Qwen2.5-1.5B-Instruct-Q4_K_M.gguf --local-dir .
#  huggingface-cli download bartowski/Qwen2.5-1.5B-Instruct-GGUF Qwen2.5-3B-Instruct-Q4_K_M.gguf --local-dir .
#
# Deuxième possibilité : au lieu de la commande huggingface_cli, exécuter le script python ci-dessous :
# from huggingface_hub import hf_hub_download
# token = "hf_..."                                  # indiquer ici le token à utiliser pour le téléchargement
# hf_hub_download(
#     repo_id="bartowski/Qwen2.5-3B-Instruct-GGUF", # indiquer ici le modèle à télécharger
#     filename="Qwen2.5-3B-Instruct-Q4_K_M.gguf",   # indiquer ici le fichier dans lequel le modèle doit être enregistré
#     local_dir=".",
#     token=token,
#     local_dir_use_symlinks=False
# )
# print("✅ Qwen2.5-3B téléchargé !")

from llama_cpp import Llama
import textwrap
import time
import json
from datetime import datetime
import os

# ================== CONFIGURATION ==================
HISTORY_FILE = "chat_history.json"

def load_history():
    if os.path.exists(HISTORY_FILE):
        try:
            with open(HISTORY_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except:
            return []
    return []

def save_history(history):
    with open(HISTORY_FILE, "w", encoding="utf-8") as f:
        json.dump(history, f, ensure_ascii=False, indent=2)

print("=== Chat Local Qwen2.5 ===\n")

# Choix du modèle
print("1 → Qwen2.5-1.5B (rapide)")
print("3 → Qwen2.5-3B   (meilleure qualité)")
choice = input("Votre choix (1 ou 3) : ").strip()

model_path = "Qwen2.5-3B-Instruct-Q4_K_M.gguf" if choice == "3" else "Qwen2.5-1.5B-Instruct-Q4_K_M.gguf"
print(f"→ Chargement de {model_path}...\n")

# Chargement du modèle
llm = Llama(
    model_path="Models/"+model_path,
    n_gpu_layers=0,
    n_ctx=2048,
    n_threads=4,
    n_batch=512,
    n_ubatch=256,
    use_mmap=False,
    verbose=False,
    temperature=0.7,      # Valeur par défaut plus "créative"
)

print("✅ Modèle chargé ! Tape 'bye' pour quitter.\n")

history = load_history()   # Charge l'historique précédent
conversation = []          # Pour l'affichage

while True:
    user_input = input(">>> ")
    if user_input.strip().lower() in ["bye", "exit", "quit"]:
        print("Good bye!")
        save_history(history)
        break

    # Construction du prompt
    prompt = "".join([msg["content"] for msg in history]) + f"<|im_start|>user\n{user_input}<|im_end|>\n<|im_start|>assistant\n"

    start = time.time()
    output = llm(prompt, max_tokens=1024, stop=["<|im_start|>", "<|im_end|>"], temperature=0.7)
    duration = time.time() - start

    answer = output['choices'][0]['text'].strip()

    # Affichage
    print("\n" + textwrap.fill(answer, width=80))
    print(f"⏱️  Temps : {duration:.2f} secondes\n")

    # Sauvegarde dans l'historique
    history.append({"role": "user", "content": f"<|im_start|>user\n{user_input}<|im_end|>\n"})
    history.append({"role": "assistant", "content": f"<|im_start|>assistant\n{answer}<|im_end|>\n"})

    conversation.append(("User", user_input))
    conversation.append(("Assistant", answer))

# Sauvegarde à la fin
save_history(history)
print("Historique sauvegardé dans chat_history.json")