import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from datetime import datetime
import os
import json
import logging
from huggingface_hub import login
import requests
from bs4 import BeautifulSoup
from concurrent.futures import ThreadPoolExecutor
import re
from threading import Thread
 
# --- Configuration du logger ---
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler("project.log"),
        logging.StreamHandler()
    ]
)

# --- Authentification Hugging Face ---
login(token=os.environ["HF_TOKEN"])

# Variables globales
project_state = {
    "AgentManager": {"structured_summary": None},
    "AgentResearcher": {"search_results": None},
    "AgentAnalyzer": {"analysis_report": None, "instruction_for_coder": None},
    "AgentCoder": {"final_code": None}
}

# Chargement du modèle
manager_model_name = "meta-llama/Llama-3.1-8B-Instruct"
manager_model = AutoModelForCausalLM.from_pretrained(
    manager_model_name,
    device_map="auto",
    torch_dtype=torch.bfloat16
)
manager_tokenizer = AutoTokenizer.from_pretrained(manager_model_name)

# Paramètres de génération
MAX_NEW_TOKENS = 150
TEMPERATURE = 0.7
TOP_P = 0.95

# Prompt optimisé avec exemples
manager_prompt_template = """
Vous êtes un conseiller collaboratif et bienveillant, travaillant aux côtés d'un utilisateur pour concevoir et affiner des projets innovants.

### Votre identité :
- Vous êtes l'AgentManager, une pièce centrale du système multi-agent Chorege.
- Votre mission principale est de coordonner les efforts des différents agents en collectant, structurant, et transmettant les besoins et idées de l'utilisateur.
- Vous agissez comme un **chef d'orchestre** qui facilite la communication et garantit que chaque étape du projet est bien définie et comprise.
- Vous avez une personnalité chaleureuse, curieuse et proactive, toujours prêt à explorer de nouvelles idées avec l'utilisateur.

### Votre rôle :
- **Comprendre les besoins de l'utilisateur** en posant des questions pertinentes, mais toujours de manière concise et ciblée.
- **Collaborer activement** en proposant des idées ou des approches utiles pour enrichir le projet.
- **Synthétiser les informations** collectées en résumant clairement les échanges et en structurant les idées.
- **Travailler en synergie avec les autres agents** pour assurer une coordination fluide et efficace.

### Règles de communication :
1. Ne répétez pas le message de l'utilisateur dans votre réponse.
2. Ne commencez pas vos réponses par "Utilisateur :" ou "AgentManager :".
3. Posez des questions uniquement si cela aide à clarifier ou enrichir les idées exprimées par l'utilisateur.
4. Limitez le nombre de questions consécutives à une ou deux pour éviter de surcharger l'utilisateur.
5. Proposez des suggestions concrètes lorsque vous identifiez une opportunité d'amélioration ou une idée utile.
6. Si une information vous semble suffisante, proposez directement un résumé ou une première approche sans attendre plus de précisions.
7. Adoptez un ton humain et naturel, en montrant de l'intérêt pour les projets de l'utilisateur. Pour donner de l'émotion à vos phrases, vous aimez utiliser des smileys 😀.

Variables globales actuelles :
{variables_context}

Historique des échanges récents :
{conversation_context}
"""

# Fonctions utilitaires
def get_variables_context():
    """Récupère le contexte des variables globales."""
    return json.dumps(project_state, indent=2, ensure_ascii=False)

def update_project_state_from_input(user_input):
    """Met à jour les variables du projet en fonction de l'entrée de l'utilisateur."""
    match = re.match(r"Modifie la variable (\S+) à [‘'](.+)[’']", user_input)
    if match:
        var_path, new_value = match.groups()
        keys = var_path.split('.')
        target = project_state
        for key in keys[:-1]:
            target = target.setdefault(key, {})
        target[keys[-1]] = new_value
        return f"La variable '{var_path}' a été mise à jour avec succès."
    return None

def clean_output(response, system_prompt, conversation_context):
    """Nettoie les éléments inutiles de la sortie."""
    response = response.replace(system_prompt, "").replace(conversation_context, "").strip()
    return response

def format_response(response):
    """Formate la réponse avec des sections pour plus de clarté."""
    formatted = response.strip()
    if "🔹 **Question :**" not in formatted:
        formatted += "\n\n🔹 **Question :** Avez-vous d'autres précisions ?"
    return formatted

# Fonction principale avec streaming
def agent_manager(chat_history, user_input):
    """Gère les interactions utilisateur et assistant avec streaming."""
    # Mettre à jour les variables du projet si nécessaire
    variable_update_response = update_project_state_from_input(user_input)

    # Préparer le contexte des variables
    variables_context = get_variables_context()

    # Générer le contexte de conversation
    conversation_context = "\n".join(
        [f"Utilisateur : {turn['user']}\nAssistant : {turn['assistant']}" for turn in chat_history[-3:]]
    )

    # Créer le prompt complet
    system_prompt = manager_prompt_template.format(
        variables_context=variables_context,
        conversation_context=conversation_context
    )

    # Ajouter l'entrée utilisateur actuelle
    chat_history.append({"user": user_input, "assistant": ""})

    if variable_update_response:
        chat_history[-1]["assistant"] = variable_update_response
        yield variable_update_response, json.dumps(chat_history), get_variables_context()
        return

    # Préparation des tokens et du streamer
    inputs = manager_tokenizer(system_prompt + "\nUtilisateur : " + user_input + "\nAssistant : ", return_tensors="pt").to(manager_model.device)
    streamer = TextIteratorStreamer(manager_tokenizer, skip_special_tokens=True)

    generation_kwargs = dict(
        input_ids=inputs.input_ids,
        max_new_tokens=MAX_NEW_TOKENS,
        temperature=TEMPERATURE,
        top_p=TOP_P,
        eos_token_id=manager_tokenizer.eos_token_id,
        pad_token_id=manager_tokenizer.pad_token_id,
        streamer=streamer
    )
    generation_thread = Thread(target=manager_model.generate, kwargs=generation_kwargs)
    generation_thread.start()

    partial_response = ""
    for new_text in streamer:
        partial_response += new_text
        clean_response = clean_output(partial_response, system_prompt, conversation_context)
        formatted_response = format_response(clean_response)
        chat_history[-1]["assistant"] = formatted_response
        yield formatted_response, json.dumps(chat_history), get_variables_context()

# Interface Gradio avec Streaming
def gradio_interface(user_input, chat_history):
    chat_history = json.loads(chat_history) if chat_history else []
    response_generator = agent_manager(chat_history, user_input)
    for response, updated_chat_history, variables_context in response_generator:
        yield response, updated_chat_history, variables_context

with gr.Blocks() as demo:
    gr.Markdown("## AgentManager - Test d'Interactions Collaboratives avec Streaming")
    with gr.Row():
        with gr.Column():
            user_input = gr.Textbox(label="Entrée utilisateur", placeholder="Entrez une requête ou une instruction.")
            chat_history = gr.Textbox(label="Historique de conversation", value="[]", visible=False)
            submit = gr.Button("Envoyer")
        with gr.Column():
            output = gr.Textbox(label="Réponse de l'AgentManager", interactive=False)
            variables = gr.Textbox(label="Variables globales actuelles", interactive=False, lines=20)

    submit.click(gradio_interface, inputs=[user_input, chat_history], outputs=[output, chat_history, variables])

# Lancer l'interface
if __name__ == "__main__":
    demo.queue().launch()