import streamlit as st
import os
import glob
import shutil
from workflow_Agent import app  # Ton graphe compilé
from langchain_core.messages import HumanMessage
import pandas as pd

# --- CONFIGURATION ET NETTOYAGE AU DÉMARRAGE ---
st.set_page_config(page_title="Dataltist AI Assistant", layout="wide")
st.title("📊 Dataltist BI Chat")

from langfuse.langchain import CallbackHandler
 
# Initialize Langfuse CallbackHandler for Langchain (tracing)
langfuse_handler = CallbackHandler()


# Initialisation du dossier de travail
if "app_initialized" not in st.session_state:
    if os.path.exists("outputs"):
        shutil.rmtree("outputs")
    os.makedirs("outputs", exist_ok=True)
    if not os.path.exists("data"):
        os.makedirs("data", exist_ok=True)
    st.session_state.app_initialized = True
    st.session_state.query_count = 0
    st.session_state.messages = []

# --- 1. SIDEBAR : GESTION DES DONNÉES ---
file_path = None

with st.sidebar:
    st.header("📁 Données")
    uploaded_file = st.file_uploader("Charge ton fichier (CSV/Excel)", type=["csv", "xlsx"])
    
    if uploaded_file:
        # 1. Définition et sauvegarde du fichier
        file_path = os.path.join("data", uploaded_file.name)
        with open(file_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        st.success(f"Fichier '{uploaded_file.name}' prêt !") 

        df = pd.read_csv(file_path) if file_path.endswith(".csv") else pd.read_excel(file_path)
        columns_list = df.columns.tolist()

        # 2. Initialisation ou Mise à jour du Session State
        # On utilise le nom exact de ta classe : AgentState
        if "AgentState" not in st.session_state:
            st.session_state.AgentState = {
                "messages": [], # Liste vide pour l'historique
                "current_df_path": file_path, # Le chemin vers le fichier uploadé
                "Data_colomns" : columns_list, 
                "generated_charts": [] # Liste vide pour les futurs PNG
            }
        else:
            # Si l'utilisateur change de fichier, on met à jour le chemin
            st.session_state.AgentState["current_df_path"] = file_path

  
    st.write("---")
    if st.button("🗑️ Effacer la discussion & fichiers", use_container_width=True):
        st.session_state.messages = []
        st.session_state.query_count = 0
        if os.path.exists("outputs"):
            shutil.rmtree("outputs")
        os.makedirs("outputs", exist_ok=True)
        st.rerun()

# --- 2. AFFICHAGE DE L'HISTORIQUE ---
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        # Réaffichage des graphiques archivés pour ce message
        if "charts" in message and message["charts"]:
            cols = st.columns(min(len(message["charts"]), 2))
            for idx, path in enumerate(message["charts"]):
                if os.path.exists(path):
                    cols[idx % 2].image(path, use_container_width=True)

# --- 3. INTERACTION AVEC LES AGENTS ---
if prompt := st.chat_input("Pose ta question sur tes données..."):
    if not file_path:
        st.error("Veuillez d'abord charger un fichier dans la barre latérale.")
        st.stop()

    # Archivage : Création du dossier spécifique pour cette requête
    current_query_id = st.session_state.query_count
    query_folder = os.path.join("outputs", f"query_{current_query_id}")
    os.makedirs(query_folder, exist_ok=True)

    # Affichage message utilisateur
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Appel des agents
    with st.chat_message("assistant"):
        with st.status("Analyse Dataltist en cours...", expanded=True) as status:
            
            inputs = {
                "messages": [HumanMessage(content=prompt)],
                "current_df_path": file_path
            }
            
            # Exécution du graphe
            final_state = app.invoke(
                inputs, 
                config={"callbacks": [langfuse_handler], "run_name": f"Query_{st.session_state.query_count}"}
            )
            
            # Récupération de la réponse finale (Reporter)
            response_text = final_state["messages"][-1].content
            st.markdown(response_text)
            
            # --- GESTION DES FICHIERS GÉNÉRÉS ---
            new_files_paths = []
            # On scanne la racine de 'outputs/' pour trouver ce que l'Exécuteur a créé
            raw_files = [f for f in glob.glob("outputs/*") if os.path.isfile(f)]
            
            if raw_files:
                st.write("---")
                st.subheader("📊 Résultats de l'analyse")
                
                # On déplace les fichiers vers le dossier de la query pour l'historique
                for f in raw_files:
                    dest_path = os.path.join(query_folder, os.path.basename(f))
                    shutil.move(f, dest_path)
                    new_files_paths.append(dest_path)

                # Affichage des images et fichiers Excel
                images = [p for p in new_files_paths if p.lower().endswith(('.png', '.jpg'))]
                docs = [p for p in new_files_paths if p.lower().endswith(('.xlsx', '.csv'))]

                if images:
                    cols = st.columns(min(len(images), 2))
                    for idx, img_path in enumerate(images):
                        cols[idx % 2].image(img_path, use_container_width=True, caption=f"Visuel {idx+1}")
                
                if docs:
                    for doc_path in docs:
                        with open(doc_path, "rb") as f:
                            st.download_button(
                                label=f"📥 Télécharger {os.path.basename(doc_path)}",
                                data=f,
                                file_name=os.path.basename(doc_path),
                                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                            )
            
            status.update(label="Analyse terminée !", state="complete")

    # Mise à jour de la session et du compteur
    st.session_state.messages.append({
        "role": "assistant", 
        "content": response_text, 
        "charts": new_files_paths
    })
    st.session_state.query_count += 1