#================imports==============

import uuid
import requests
import os
os.environ["USER_AGENT"] = "RAG-App/1.0"
from typing import Dict, List, Any, Generator
from dotenv import load_dotenv

from bs4 import BeautifulSoup
from langchain_core.globals import set_llm_cache
from langchain_core.caches import InMemoryCache

from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Weaviate
from langchain_community.vectorstores import FAISS
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_classic.chains import create_retrieval_chain

from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory

import gradio as gr

#================== CONFIG==================

load_dotenv()

set_llm_cache(InMemoryCache())
api_key = os.environ["GROQ_API_KEY"]
print("api chargée:" if api_key else "y'a probleme!!")

#========== charger et decouper documents=================

urls = [
    "https://fr.wikipedia.org/wiki/%C3%89levage",
    "https://fr.wikipedia.org/wiki/La_P%C3%AAche"
]

loader = WebBaseLoader(urls,
                      requests_kwargs={
                          "headers": {
                              "User-Agent": "RAG-App/1.0"
                          }
                      }
)
docs = loader.load()

splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
chunks = splitter.split_documents(docs)

#============embeding et indexation vers faiss_db================

embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

faiss_db = FAISS.from_documents(
    documents=chunks,
    embedding=embeddings
)

retriever = faiss_db.as_retriever(search_type="similarity", search_kwargs={"k": 3})

#=============== LLM et Prompt=================

llm = ChatGroq(
    model="llama-3.3-70b-versatile",
    temperature=0.0,
    max_tokens=1200,
    streaming=True  # Activer le streaming pour une réponse en temps réel
)

prompt = ChatPromptTemplate.from_messages([
    ("system", """Tu es un assistant expert en dans le domaine de l'elevage et la pêche. Réponds clairement.
Si tu ne connais pas, n'invente pas. Garde un ton amical.
Contexte :
{context}"""),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
])

#============= CHAINE DE RECUPERATION=======

stuff_chain = create_stuff_documents_chain(llm, prompt)
rag_chain = create_retrieval_chain(retriever, stuff_chain)

# ====== GESTION DE L'HISTORIQUE ======

store = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]

# ====== CHAIN AVEC MÉMOIRE ======

convers_chain = RunnableWithMessageHistory(
    rag_chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer"
)

SESSION_ID = str(uuid.uuid4())  # session globale

# ================= CSS PERSONNALISÉ POUR LE STYLE CHATGPT ====================

custom_css = """
/* Style global */
.gradio-container {
    max-width: 100% !important;
    margin: 0 !important;
    padding: 0 !important;
}

/* Style de la sidebar */
.sidebar {
    background: #202123;
    color: white;
    height: 100vh;
    overflow-y: auto;
}

.sidebar-header {
    padding: 12px 16px;
    border-bottom: 1px solid #4d4d4f;
    margin-bottom: 8px;
}

.new-chat-btn {
    width: 100%;
    padding: 8px 16px;
    background: transparent;
    border: 1px solid #4d4d4f;
    color: white;
    border-radius: 8px;
    cursor: pointer;
    font-size: 14px;
    text-align: left;
    transition: background 0.3s;
}

.new-chat-btn:hover {
    background: #2b2c2f;
}

.history-item {
    padding: 12px 16px;
    cursor: pointer;
    border-radius: 8px;
    margin: 4px 8px;
    transition: background 0.3s;
    word-wrap: break-word;
    color: #ececec;
}

.history-item:hover {
    background: #2b2c2f;
}

.history-item.active {
    background: #343541;
}

.main-chat-area {
    background: #343541;
    height: 100vh;
}

.chatbot-container {
    height: calc(100vh - 180px) !important;
}

.chatbot-container > div {
    border: none !important;
    background: #343541 !important;
}
"""

# ================= INTERFACE GRADIO STYLE CHATGPT ====================

with gr.Blocks(css=custom_css, theme="soft") as demo:
    with gr.Row(equal_height=True):
        # ==================== SIDEBAR GAUCHE (Style ChatGPT) ====================
        with gr.Column(scale=1, min_width=260, elem_classes="sidebar"):
            gr.HTML("""
                <div class="sidebar-header">
                    <h2 style="color: white; font-size: 16px; margin: 0;">🐟 RAG Assistant</h2>
                    <p style="color: #8e8ea0; font-size: 12px; margin: 4px 0 0 0;">Élevage & Pêche</p>
                </div>
            """)
            
            # Bouton Nouvelle conversation
            new_chat_btn = gr.Button("➕ Nouvelle conversation", elem_classes="new-chat-btn")
            
            gr.HTML('<div style="padding: 8px 16px; color: #8e8ea0; font-size: 12px; margin-top: 16px;">HISTORIQUE</div>')
            
            # Liste des conversations précédentes
            history_radio = gr.Radio(
                choices=[],
                label=None,
                interactive=True,
                elem_classes="history-radio",
                container=False
            )
            
            # Variables d'état pour stocker les conversations
            conversations_state = gr.State([])
            current_conversation_id = gr.State("")
            
            # Message pour confirmer les actions
            status_msg = gr.HTML(visible=False)
        
        # ==================== ZONE PRINCIPALE DE CHAT (Style ChatGPT) ====================
        with gr.Column(scale=4, elem_classes="main-chat-area"):
            # Header
            gr.HTML("""
                <div style="padding: 16px; border-bottom: 1px solid #4d4d4f; background: #343541;">
                    <h2 style="color: white; font-size: 18px; margin: 0; text-align: center;">
                        Assistant RAG - Élevage & Pêche
                    </h2>
                </div>
            """)
            
            # Chatbot
            chatbot = gr.Chatbot(
                label=None,
                height=500,
                elem_classes="chatbot-container",
                show_label=False,
                bubble_full_width=False,
                avatar_images=(None, "🐟")
            )
            
            # Zone de saisie style ChatGPT
            with gr.Row():
                msg_input = gr.Textbox(
                    show_label=False,
                    placeholder="Envoyez un message...",
                    scale=9,
                    container=False,
                    lines=1,
                    max_lines=5
                )
                send_btn = gr.Button("↑", variant="primary", scale=1, min_width=40)
            
            # Texte de copyright en bas
            gr.HTML("""
                <div style="text-align: center; padding: 8px; color: #8e8ea0; font-size: 11px;">
                    RAG Assistant peut faire des erreurs. Vérifiez les informations importantes.
                </div>
            """)
    
    # ==================== FONCTIONS DE GESTION ====================
    
    def create_new_conversation(conversations):
        """Crée une nouvelle conversation et retourne l'ID."""
        conversation_id = str(uuid.uuid4())
        title = "Nouvelle conversation"
        conversations.append({"id": conversation_id, "title": title, "messages": []})
        return conversations, conversation_id, title, gr.update(choices=[c["title"] for c in conversations], value=None)
    
    def load_conversation(selected_title, conversations):
        """Charge une conversation existante."""
        if not selected_title:
            return [], "", ""
        
        for conv in conversations:
            if conv["title"] == selected_title:
                chat_history = []
                for msg in conv["messages"]:
                    if msg["role"] == "user":
                        chat_history.append((msg["content"], None))
                    else:
                        if chat_history:
                            chat_history[-1] = (chat_history[-1][0], msg["content"])
                        else:
                            chat_history.append((None, msg["content"]))
                return chat_history, conv["id"], conv["title"]
        
        return [], "", ""
    
    def send_message(message, chat_history, conversations, current_conv_id, chat_state):
        """Envoie un message et met à jour la conversation."""
        if not message or not message.strip():
            return "", chat_history, conversations, current_conv_id, chat_state
        
        # Créer une nouvelle conversation si nécessaire
        if not current_conv_id:
            current_conv_id = str(uuid.uuid4())
            # Utiliser les premiers mots comme titre
            title = message[:50] + "..." if len(message) > 50 else message
            conversations.append({
                "id": current_conv_id, 
                "title": title, 
                "messages": [{"role": "user", "content": message}]
            })
        else:
            # Ajouter le message à la conversation existante
            for conv in conversations:
                if conv["id"] == current_conv_id:
                    conv["messages"].append({"role": "user", "content": message})
                    if conv["title"] == "Nouvelle conversation":
                        conv["title"] = message[:50] + "..." if len(message) > 50 else message
                    break
        
        # Ajouter le message utilisateur au chatbot
        chat_history.append((message, None))
        
        # Obtenir la réponse de l'assistant
        result = convers_chain.invoke(
            {"input": message},
            config={"configurable": {"session_id": SESSION_ID}}
        )
        
        response = result.get("answer", str(result))
        
        # Ajouter la réponse à la conversation
        for conv in conversations:
            if conv["id"] == current_conv_id:
                conv["messages"].append({"role": "assistant", "content": response})
                break
        
        # Mettre à jour le chatbot
        chat_history[-1] = (message, response)
        
        # Mettre à jour les choix du radio
        choices = [conv["title"] for conv in conversations]
        
        return "", chat_history, conversations, current_conv_id, gr.update(choices=choices, value=conversations[-1]["title"] if conversations else None)
    
    def clear_chat():
        """Efface le chat et commence une nouvelle conversation."""
        return [], "", None, gr.update(choices=[], value=None)
    
    # ==================== GESTIONNAIRES D'ÉVÉNEMENTS ====================
    
    # Envoi de message
    msg_input.submit(
        send_message,
        inputs=[msg_input, chatbot, conversations_state, current_conversation_id, gr.State([])],
        outputs=[msg_input, chatbot, conversations_state, current_conversation_id, history_radio]
    ).then(
        lambda: gr.update(value=""),
        outputs=[msg_input]
    )
    
    send_btn.click(
        send_message,
        inputs=[msg_input, chatbot, conversations_state, current_conversation_id, gr.State([])],
        outputs=[msg_input, chatbot, conversations_state, current_conversation_id, history_radio]
    ).then(
        lambda: gr.update(value=""),
        outputs=[msg_input]
    )
    
    # Nouvelle conversation
    new_chat_btn.click(
        create_new_conversation,
        inputs=[conversations_state],
        outputs=[conversations_state, current_conversation_id, gr.State(""), history_radio]
    ).then(
        clear_chat,
        outputs=[chatbot, current_conversation_id, history_radio, history_radio]
    )
    
    # Charger une conversation depuis l'historique
    history_radio.change(
        load_conversation,
        inputs=[history_radio, conversations_state],
        outputs=[chatbot, current_conversation_id, gr.State("")]
    )
    
    # Support de la touche Entrée (Shift+Entrée pour nouvelle ligne)
    def handle_enter_key(text, event: gr.EventData):
        if event.key == "Enter" and not event.shift:
            return text, gr.update()
        return text, gr.update()
    
    msg_input.key_up(
        handle_enter_key,
        inputs=[msg_input],
        outputs=[msg_input, chatbot]
    )

# ===================LANCEMENT ================

demo.launch(share=False)