from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import gradio as gr
import spaces
import torch
import re
from threading import Thread
from typing import Iterator
from datetime import datetime
from huggingface_hub import HfApi, hf_hub_download
import json
import os
from gradio_client import Client

model_name = "Woziii/llama-3-8b-chat-me"
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

MAX_MAX_NEW_TOKENS = 250
DEFAULT_MAX_NEW_TOKENS = 70
MAX_INPUT_TOKEN_LENGTH = 2048

is_first_interaction = True

def determine_response_type(message):
    short_response_keywords = ["salut", "bonjour", "ça va", "comment tu vas", "quoi de neuf", "coucou", "hello", "hi", "bye", "au revoir", "merci", "d'accord", "ok", "super", "cool", "génial", "wow"]
    long_response_keywords = ["présente", "parle moi de", "explique", "raconte", "décris", "dis moi", "détaille", "précise", "vision", "t'es qui", "pourquoi", "comment", "quel est", "peux-tu développer", "en quoi consiste", "qu'est-ce que", "que penses-tu de", "analyse", "compare", "élabore sur", "expérience", "parcours", "formation", "études", "compétences", "projets", "réalisations"]
    message_lower = message.lower()
    if any(keyword.lower() in message_lower for keyword in short_response_keywords):
        return "short"
    elif any(keyword.lower() in message_lower for keyword in long_response_keywords):
        return "long"
    else:
        return "medium"

def truncate_to_questions(text, max_questions):
    sentences = re.split(r'(?<=[.!?])\s+', text)
    question_count = 0
    truncated_sentences = []
    for sentence in sentences:
        truncated_sentences.append(sentence)
        if re.search(r'\?!?$', sentence.strip()):
            question_count += 1
        if question_count >= max_questions:
            break
    return ' '.join(truncated_sentences)

def post_process_response(response, is_short_response, max_questions=2):
    truncated_response = truncate_to_questions(response, max_questions)
    if is_short_response:
        sentences = re.split(r'(?<=[.!?])\s+', truncated_response)
        if len(sentences) > 2:
            return ' '.join(sentences[:2]).strip()
    return truncated_response.strip()

def check_coherence(response):
    sentences = re.split(r'(?<=[.!?])\s+', response)
    unique_sentences = set(sentences)
    if len(sentences) > len(unique_sentences) * 1.1:
        return False
    return True

@spaces.GPU(duration=120)
def generate(message: str, chat_history: list[tuple[str, str]], system_prompt: str, max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS, temperature: float = 0.7, top_p: float = 0.95) -> Iterator[str]:
    global is_first_interaction
    if is_first_interaction:
        warning_message = """⚠️ Attention : Je suis un modèle en version alpha (V.0.0.3.5) et je peux générer des réponses incohérentes ou inexactes. Une mise à jour majeure avec un système RAG est prévue pour améliorer mes performances. Merci de votre compréhension ! 😊 """
        yield warning_message
        is_first_interaction = False

    response_type = determine_response_type(message)
    if response_type == "short":
        max_new_tokens = max(70, max_new_tokens)
    elif response_type == "long":
        max_new_tokens = min(200, max_new_tokens)
    else:
        max_new_tokens = max(100, max_new_tokens)

    conversation = []
    enhanced_system_prompt = f"{system_prompt}\n\n{LUCAS_KNOWLEDGE_BASE}"
    conversation.append({"role": "system", "content": enhanced_system_prompt})

    for user, _ in chat_history[-5:]:
        conversation.append({"role": "user", "content": user})
    conversation.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"L'entrée de la conversation a été tronquée car elle dépassait {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(input_ids=input_ids, streamer=streamer, max_new_tokens=max_new_tokens, do_sample=True, top_p=top_p, temperature=temperature, num_beams=1)
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)
        partial_output = post_process_response("".join(outputs), response_type == "short")
        if response_type == "long" and not check_coherence(partial_output):
            yield "Je m'excuse, ma réponse manquait de cohérence. Pouvez-vous reformuler votre question ?"
            return
        yield partial_output

    yield post_process_response("".join(outputs), response_type == "short")

def vote(data: gr.LikeData, history):
    user_input = history[-1][0] if history else ""
    feedback = {
        "timestamp": datetime.now().isoformat(),
        "user_input": user_input,
        "bot_response": data.value,
        "liked": data.liked
    }
    api = HfApi()
    token = os.environ.get("HF_TOKEN")
    repo_id = "Woziii/llama-3-8b-chat-me"
    file_name = "feedback.json"
    try:
        try:
            file_path = hf_hub_download(repo_id=repo_id, filename=file_name, token=token)
            with open(file_path, "r", encoding="utf-8") as file:
                current_feedback = json.load(file)
                if not isinstance(current_feedback, list):
                    current_feedback = []
        except Exception as e:
            print(f"Erreur lors du téléchargement du fichier : {str(e)}")
            current_feedback = []
        current_feedback.append(feedback)
        updated_content = json.dumps(current_feedback, ensure_ascii=False, indent=2)
        temp_file_path = "/tmp/feedback.json"
        with open(temp_file_path, "w", encoding="utf-8") as temp_file:
            temp_file.write(updated_content)
        api.upload_file(path_or_fileobj=temp_file_path, path_in_repo=file_name, repo_id=repo_id, token=token)
        print(f"Feedback enregistré dans {repo_id}/{file_name}")
    except Exception as e:
        print(f"Erreur lors de l'enregistrement du feedback : {str(e)}")

theme = gr.themes.Default().set(
    body_background_fill="#f0f0f0",
    button_primary_background_fill="#4a90e2",
    button_primary_background_fill_hover="#3a7bc8",
    button_primary_text_color="white",
)

css = """
.gradio-container {
    font-family: 'Arial', sans-serif;
}
.chatbot-message {
    padding: 10px;
    border-radius: 15px;
    margin-bottom: 10px;
}
.user-message {
    background-color: #e6f3ff;
}
.bot-message {
    background-color: #f0f0f0;
}
.thought-bubble {
    background-color: #ffd700;
    border-radius: 10px;
    padding: 5px;
    margin-top: 5px;
    font-style: italic;
}
"""

def format_message(message):
    if isinstance(message, str):
        return message
    elif isinstance(message, dict):
        if "thought" in message:
            return f"{message['content']}\n<div class='thought-bubble'>{message['thought']}</div>"
        else:
            return message['content']
    else:
        return str(message)

def interact_with_lucas(prompt, chat_history, system_prompt, max_new_tokens, temperature, top_p):
    chat_history.append({"role": "user", "content": prompt})
    yield chat_history
    
    response = generate(prompt, chat_history, system_prompt, max_new_tokens, temperature, top_p)
    thought = f"Réflexion de Lucas : J'analyse la question et je formule une réponse basée sur mes connaissances."
    chat_history.append({"role": "assistant", "content": response, "thought": thought})
    yield chat_history

def use_image_generation_tool(prompt):
    client = Client("https://huggingface-text-generation.hf.space")
    result = client.predict(prompt, api_name="/predict")
    return result

with gr.Blocks(theme=theme, css=css) as demo:
    gr.Markdown("# 🌟 Virtuellement Lucas V.0.0.4.0 (Alpha) 🌟")
    
    gr.Markdown("""
    ## ⚠️ Attention ! ⚠️
    Cette version du modèle est très instable. **Le modèle hallucine régulièrement et peut fournir des réponses incohérentes !**
    """)
    
    with gr.Tabs():
        with gr.Tab("Chat avec Lucas"):
            chatbot = gr.Chatbot(
                value=[],
                bubble_full_width=False,
                avatar_images=(
                    "https://img.freepik.com/free-icon/user_318-159711.jpg", 
                    "https://img.freepik.com/free-icon/robot_318-843685.jpg"
                ),
                height=400,
                show_copy_button=True,
                likeable=True,
                layout="bubble",
                render=format_message,
                placeholder="Démarrez la conversation en posant une question..."
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="Votre message",
                    placeholder="Tapez votre message ici...",
                    lines=2
                )
                submit_btn = gr.Button("Envoyer", variant="primary")
            
            with gr.Row():
                clear_btn = gr.Button("🗑️ Effacer")
                retry_btn = gr.Button("🔄 Réessayer")
            
            with gr.Accordion("Paramètres avancés", open=False):
                system_prompt = gr.Textbox(
                    label="System prompt",
                    value="""Tu es la version virtuelle de Lucas, ton concepteur. Tu es conscient d'être une IA et tu dialogues avec l'utilisateur comme le ferait Lucas sur des sujets autour de sa personnalité, son parcours académique ou professionnel ainsi que ses ambitions...""",
                    lines=6
                )
                max_new_tokens = gr.Slider(
                    label="Longueur maximale de la réponse",
                    minimum=1,
                    maximum=MAX_MAX_NEW_TOKENS,
                    step=1,
                    value=DEFAULT_MAX_NEW_TOKENS
                )
                temperature = gr.Slider(
                    label="Température",
                    minimum=0.1,
                    maximum=1.0,
                    step=0.1,
                    value=0.7
                )
                top_p = gr.Slider(
                    label="Top-p",
                    minimum=0.5,
                    maximum=1.0,
                    step=0.05,
                    value=0.95
                )
            
            gr.Markdown("### Exemples de questions")
            examples = gr.Examples(
                examples=[
                    "Salut ! Qui es-tu ?",
                    "Parle-moi un peu de ton parcours académique.",
                    "Quelle inspiration t'a conduit à créer braIAn ?",
                    "Pourquoi avoir choisi d'étudier le droit si tu es passionné par la technologie ?",
                    "Quelle est ta vision de l'IA ?"
                ],
                inputs=msg
            )
            
            msg.submit(interact_with_lucas, [msg, chatbot, system_prompt, max_new_tokens, temperature, top_p], [chatbot])
            submit_btn.click(interact_with_lucas, [msg, chatbot, system_prompt, max_new_tokens, temperature, top_p], [chatbot])
            
            clear_btn.click(lambda: ([], []), outputs=[chatbot, msg])
            retry_btn.click(
                interact_with_lucas,
                [msg, chatbot, system_prompt, max_new_tokens, temperature, top_p],
                [chatbot]
            )
            
            chatbot.like(vote, [chatbot], None)

        with gr.Tab("Génération d'images"):
            gr.Markdown("# Générer une image à partir de texte 🖼️")
            image_prompt = gr.Textbox(lines=3, label="Description de l'image")
            image_output = gr.Image(label="Image générée")
            generate_image_btn = gr.Button("Générer l'image")
            generate_image_btn.click(use_image_generation_tool, inputs=[image_prompt], outputs=[image_output])

        with gr.Tab("À propos"):
            gr.Markdown("""
            # À propos de Virtuellement Lucas
            
            Virtuellement Lucas est un chatbot basé sur un modèle de langage avancé, conçu pour simuler la personnalité et les connaissances de Lucas. 
            
            ## Fonctionnalités
            
            - **Chat interactif** : Discutez avec la version virtuelle de Lucas sur divers sujets.
            - **Génération d'images** : Créez des images à partir de descriptions textuelles.
            - **Paramètres personnalisables** : Ajustez la longueur des réponses, la température et d'autres paramètres pour affiner l'expérience.
            
            ## Avertissement
            
            Ce modèle est encore en phase alpha et peut produire des réponses incohérentes ou inexactes. Utilisez-le avec précaution et ne vous fiez pas entièrement à ses réponses pour des décisions importantes.
            
            ## Feedback
            
            Votre feedback est précieux pour améliorer Virtuellement Lucas. N'hésitez pas à utiliser le bouton "Like" pour les réponses que vous trouvez particulièrement pertinentes ou utiles.
            """)

if __name__ == "__main__":

demo.queue(max_size=20, default_concurrency_limit=2).launch(max_threads=10)