import gradio as gr
from langchain_core.messages import HumanMessage, AIMessage
from src.agents.main_agent import multi_agent_graph
from src.agents.rag_agent import invocation_state
from src.configs.config import LOG_DIR
import logging
import os

LOG_FILE = os.path.join(LOG_DIR, "Agents.log")
logging.basicConfig(
    filename=LOG_FILE,
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

import pandas as pd

invocation_state.invocation_type = 'chatbot'

class GradioChatbotApp:
    def __init__(self):
        #graph_img = multi_agent_graph.get_graph(xray=True).draw_mermaid_png()
        #os.makedirs("docs", exist_ok=True)
        #with open("docs/full_graph.png", "wb") as f:
        #    f.write(graph_img)
        pass

    async def process_chat(self, history, user_input):
        messages = []
        for user, bot in history:
            messages.append(HumanMessage(content=user))
            messages.append(AIMessage(content=bot))
        messages.append(HumanMessage(content=user_input))

        logging.info(f"User query: {user_input}")
        config = {"configurable": {"thread_id": "1"}}
        result = await multi_agent_graph.ainvoke({"messages": messages}, config=config)
        bot_msg = result["messages"][-1].content
        history.append((user_input, bot_msg))
        return history, ""

    def create_interface(self):
        # Load CSV and keep only needed columns
        csv_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../dataset/docs_metadata.csv"))
        if os.path.exists(csv_path):
            df_full = pd.read_csv(csv_path)
            df = df_full[["Catégorie", "Nom du document", "Lien"]]
        else:
            df = pd.DataFrame(columns=["Catégorie", "Nom du document", "Lien"])

        # Absolute path for logo
        logo_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "docs", "logo.jpg"))

        with gr.Blocks(title="HOLOKIA RAG/MCP", theme=gr.themes.Soft()) as demo:
            with gr.Tab("Chatbot"):
                # Centered logo
                gr.HTML(f"""
                <div style='display:flex; justify-content:center; align-items:center; margin-top:30px; margin-bottom:30px;'>
                    <img src="https://holokia.com/wp-content/uploads/2025/04/cropped-logo_holokia_noir.png" style="width:130px; height:auto;" alt="Company Logo">
                </div>
                """)

                gr.Markdown("""
                ## Chatbot Juridique

                Cette plateforme d’assistance juridique intelligente est dédiée à l’analyse et à l’interprétation du droit marocain ainsi qu’aux travaux parlementaires.

                Elle s’appuie sur un système de génération augmentée par récupération (RAG), permettant de délivrer des réponses précises et documentées à partir des textes législatifs et des archives parlementaires.
                """)

                chatbot = gr.Chatbot(label="Assistant", type="tuples")
                user_input = gr.Textbox(
                    label="Votre message",
                    placeholder="Posez une question sur la loi marocaine, le parlement, ou des documents légaux...",
                    scale=7
                )

                with gr.Row():
                    gr.Examples(
                        examples=[
                            ["ما هي العقوبات المقررة لبيع المشروبات الكحولية بدون رخصة؟"],
                            ["ما هو دور البرلمان حسب الدستور المغربي 2011 ؟"],
                            ["Quel est le role du parlement dans la Constitution marocaine 2011?"],
                            ["Comment est defini le droit de greve dans la Constitution ?"]
                        ],
                        inputs=user_input
                    )

                with gr.Row():
                    send_btn = gr.Button("Envoyer", variant="primary")
                    clear_btn = gr.Button("Effacer la conversation")

                state = gr.State([])

                send_btn.click(self.process_chat, [chatbot, user_input], [chatbot, user_input])
                user_input.submit(self.process_chat, [chatbot, user_input], [chatbot, user_input])
                clear_btn.click(fn=lambda: [], inputs=None, outputs=chatbot)

            with gr.Tab("Documents"):
                gr.Markdown("## Liste des documents utilisés dans le RAG")
                gr.DataFrame(df, label="Documents", interactive=False)

        return demo

    def launch(self):
        logging.info("Launching Gradio app...")
        demo = self.create_interface()
        demo.launch(server_name="0.0.0.0", server_port=7860)

if __name__ == "__main__":
    app = GradioChatbotApp()
    app.launch()