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

import uuid
import requests
import os
os.environ["USER_AGENT"] = "RAG-App/1.0"
from typing import Dict, List, Any
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 FAISS
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.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.get("GROQ_API_KEY")
if not api_key:
    raise ValueError("❌ GROQ_API_KEY non trouvée!")
print("✅ API chargée avec succès")

#========== charger et découper documents=================

print("📥 Chargement des documents...")
urls = [
    "https://fr.wikipedia.org/wiki/%C3%89levage",
    "https://fr.wikipedia.org/wiki/La_P%C3%AAche"
]

try:
    loader = WebBaseLoader(
        urls,
        requests_kwargs={"headers": {"User-Agent": "RAG-App/1.0"}}
    )
    docs = loader.load()
    print(f"✅ {len(docs)} documents chargés")
except Exception as e:
    print(f"⚠️ Erreur de chargement: {e}")
    from langchain_core.documents import Document
    docs = [
        Document(page_content="L'élevage est l'ensemble des activités qui assurent la multiplication et l'entretien des animaux domestiques pour la production de biens et services."),
        Document(page_content="La pêche est l'activité consistant à capturer des animaux aquatiques dans leur milieu naturel.")
    ]

splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
chunks = splitter.split_documents(docs)
print(f"✅ {len(chunks)} segments créés")

#============embedding et indexation================

print("🔧 Création des embeddings...")
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={'device': 'cpu'}
)

faiss_db = FAISS.from_documents(documents=chunks, embedding=embeddings)
print("✅ Base FAISS créée")

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

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

print("🤖 Initialisation du LLM...")
llm = ChatGroq(
    model="llama-3.3-70b-versatile",
    temperature=0.0,
    max_tokens=1200
)

prompt = ChatPromptTemplate.from_messages([
    ("system", """Tu es un assistant expert en élevage et pêche. 
Réponds de manière claire et concise en français.
Si tu ne connais pas la réponse, dis-le honnêtement.
Contexte : {context}"""),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
])

#============= CHAINE DE RÉCUPÉRATION================

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]

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

# ================= CSS POUR LE STYLE ====================

custom_css = """
.sidebar {
    background: #202123 !important;
    min-height: 100vh;
}
.main-area {
    background: #343541 !important;
}
.chatbot-container {
    height: calc(100vh - 200px) !important;
}
"""

# ================= INTERFACE GRADIO (Gradio 6.0 compatible) ====================

with gr.Blocks() as demo:
    with gr.Row(equal_height=True):
        # Colonne gauche : Historique
        with gr.Column(scale=1, min_width=250, elem_classes="sidebar"):
            gr.Markdown("## 📚 Historique")
            new_chat_btn = gr.Button("➕ Nouvelle conversation", variant="secondary")
            history_radio = gr.Radio(
                choices=[],
                label="Conversations",
                interactive=True
            )
            clear_btn = gr.Button("🗑️ Effacer", variant="stop", size="sm")
        
        # Colonne droite : Chat
        with gr.Column(scale=3, elem_classes="main-area"):
            gr.Markdown("# 🤖 Assistant Élevage & Pêche")
            
            chatbot = gr.Chatbot(
                label="",
                height=500,
                show_label=False,
                avatar_images=(None, "🐟")
            )
            
            with gr.Row():
                msg_input = gr.Textbox(
                    placeholder="Posez votre question sur l'élevage ou la pêche...",
                    show_label=False,
                    scale=9,
                    container=False
                )
                send_btn = gr.Button("📤", variant="primary", scale=1)
            
            gr.Examples(
                examples=["C'est quoi l'élevage ?", "Explique la pêche", "Différence entre élevage et pêche ?"],
                inputs=msg_input
            )
    
    # États
    conversations_state = gr.State([])
    current_session_id = gr.State(str(uuid.uuid4()))
    
    # ================= FONCTIONS =================
    
    def create_new_chat(conversations):
        """Nouvelle conversation"""
        new_id = str(uuid.uuid4())
        conversations.append({"id": new_id, "title": "Nouveau chat", "messages": []})
        choices = [c["title"] for c in conversations]
        return conversations, new_id, [], gr.update(choices=choices, value=None)
    
    def load_chat(selected_title, conversations):
        """Charger une conversation"""
        if not selected_title or not conversations:
            return [], ""
        
        for conv in conversations:
            if conv["title"] == selected_title:
                history = []
                user_msg = None
                for msg in conv["messages"]:
                    if msg["role"] == "user":
                        user_msg = msg["content"]
                    else:
                        if user_msg:
                            history.append([user_msg, msg["content"]])
                            user_msg = None
                return history, conv["id"]
        return [], ""
    
    def send_message(message, chat_history, conversations, session_id):
        """Envoyer un message"""
        if not message or not message.strip():
            return "", chat_history, conversations, session_id, gr.update()
        
        # Gérer la session
        if not session_id:
            session_id = str(uuid.uuid4())
        
        # Mettre à jour la conversation dans l'historique
        conv_exists = False
        for conv in conversations:
            if conv["id"] == session_id:
                conv_exists = True
                conv["messages"].append({"role": "user", "content": message})
                if conv["title"] == "Nouveau chat":
                    conv["title"] = message[:40] + "..."
                break
        
        if not conv_exists:
            conversations.append({
                "id": session_id,
                "title": message[:40] + "...",
                "messages": [{"role": "user", "content": message}]
            })
        
        # Ajouter le message au chatbot
        chat_history.append([message, None])
        
        try:
            result = convers_chain.invoke(
                {"input": message},
                config={"configurable": {"session_id": session_id}}
            )
            response = result.get("answer", "Désolé, je n'ai pas compris.")
        except Exception as e:
            response = f"❌ Erreur: {str(e)}"
        
        # Sauvegarder la réponse
        for conv in conversations:
            if conv["id"] == session_id:
                conv["messages"].append({"role": "assistant", "content": response})
                break
        
        # Mettre à jour le chatbot
        chat_history[-1] = [message, response]
        
        # Mettre à jour la liste
        choices = [c["title"] for c in conversations]
        
        return "", chat_history, conversations, session_id, gr.update(choices=choices, value=conversations[-1]["title"] if conversations else None)
    
    def clear_all():
        """Tout effacer"""
        store.clear()
        return [], [], [], gr.update(choices=[])
    
    # ================= ÉVÉNEMENTS =================
    
    msg_input.submit(
        send_message,
        inputs=[msg_input, chatbot, conversations_state, current_session_id],
        outputs=[msg_input, chatbot, conversations_state, current_session_id, history_radio]
    )
    
    send_btn.click(
        send_message,
        inputs=[msg_input, chatbot, conversations_state, current_session_id],
        outputs=[msg_input, chatbot, conversations_state, current_session_id, history_radio]
    )
    
    new_chat_btn.click(
        create_new_chat,
        inputs=[conversations_state],
        outputs=[conversations_state, current_session_id, chatbot, history_radio]
    )
    
    history_radio.change(
        load_chat,
        inputs=[history_radio, conversations_state],
        outputs=[chatbot, current_session_id]
    )
    
    clear_btn.click(
        clear_all,
        outputs=[chatbot, conversations_state, current_session_id, history_radio]
    )

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

if __name__ == "__main__":
    print("🚀 Lancement de l'application...")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        css=custom_css,
        theme="soft"
    )