"""
MODULE JASMINE AGENT - EXECUTIVE DIRECTOR EDITION (FINAL)
=========================================================
Responsabilité : Cerveau central. Analyse stratégique, interprétation et recommandations.
Niveau : C-Level (Directeur Général).
Correctifs appliqués : Gestion variable DQ, Clamping Scores, Analyse Démographique complète.
"""

import google.generativeai as genai
from groq import Groq
import json
import streamlit as st
import re
import os
import pandas as pd
import numpy as np
import traceback

# --- IMPORTATION ROBUSTE DES MODULES ---
try:
    from module.Validation_Contract import ETL_Engine
    from module.Dashboard import RealtimeMonitor, CrisisSimulator
    from module.Data_Exploration import HospitalAdvancedAnalyzer
    from module.Indicateurs_DQ import HospitalDataQualityAnalyzer, load_from_google_sheets
    from module.Features import FeatureEngineer
    from module.Forecasting import HospitalForecaster
except ImportError as e:
    st.error(f"Erreur critique d'importation des modules : {str(e)}")

# --- STYLE ---
def apply_custom_styles():
    st.markdown("""
    <style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&family=JetBrains+Mono:wght@400&display=swap');
    .user-msg { 
        background-color: rgba(20, 30, 45, 0.9); border: 1px solid #2b3b55;
        border-left: 3px solid #64B5F6; padding: 15px; margin: 10px 0; 
        border-radius: 0px; color: #e0e0e0; font-family: 'Inter', sans-serif;
    }
    .bot-msg { 
        background-color: rgba(25, 25, 25, 0.9); border: 1px solid #333;
        border-left: 3px solid #00E676; padding: 15px; margin: 10px 0; 
        border-radius: 0px; color: #e0e0e0; font-family: 'Inter', sans-serif;
    }
    .debug-block {
        font-family: 'JetBrains Mono', monospace; font-size: 0.8rem; 
        color: #8b949e; background: #0d1117; padding: 10px; 
        border: 1px solid #30363d; border-radius: 0px; margin: 5px 0 10px 0;
    }
    </style>
    """, unsafe_allow_html=True)

class JasmineAgent:
    def __init__(self, gspread_client):
        self.gspread_client = gspread_client
        self.google_key = os.environ.get("GOOGLE_API_KEY") or st.secrets.get("GOOGLE_API_KEY")
        self.groq_key = os.environ.get("GROQ_API_KEY") or st.secrets.get("GROQ_API_KEY")
        
        self.groq_client = Groq(api_key=self.groq_key) if self.groq_key else None
        if self.google_key: genai.configure(api_key=self.google_key)

        self.MODEL_CASCADE = [
            "gemini-2.0-flash-exp",
            "llama-3.3-70b-versatile",
            "gemini-2.0-flash-lite",
            "llama-3.1-8b-instant"
        ]
        
        # --- CARTE DES OUTILS ---
        self.tools_map = {
            "run_etl_pipeline": {"func": self._run_etl_wrapper, "desc": "ETL : Mise à jour des données brutes."},
            "run_feature_engineering": {"func": self._run_features_wrapper, "desc": "Features : Transformation pour l'IA."},
            "train_forecasting_models": {"func": self._train_models_wrapper, "desc": "Training : Entraînement des modèles IA."},
            "get_future_forecast": {"func": self._get_forecast_wrapper, "desc": "Prédiction : Analyse du futur (3 mois)."},
            "get_hospital_kpis": {"func": self._get_kpis_wrapper, "desc": "Cockpit : État actuel temps réel (Monitor)."},
            "analyze_patient_stats": {"func": self._get_data_exploration_wrapper, "desc": "Data Lab : Analyse historique profonde."},
            "get_data_quality_report": {"func": self._get_dq_stats_wrapper, "desc": "Audit : Fiabilité des données."}
        }

    # --- WRAPPERS (Extraction de données brutes) ---

    def _run_etl_wrapper(self, args=None):
        try:
            engine = ETL_Engine(self.gspread_client, sheet_name="bdd_mvp_faker")
            success, result = engine.run_etl()
            if success:
                st.session_state.data_explorer = result
                return {"status": "succès", "details": result.get_summary()}
            return {"status": "erreur", "message": "Échec critique du pipeline ETL."}
        except Exception as e: return {"status": "crash", "message": str(e)}

    def _run_features_wrapper(self, args=None):
        try:
            if 'data_explorer' not in st.session_state: return {"status": "erreur", "message": "Données brutes manquantes. Lancez l'ETL."}
            engineer = FeatureEngineer(st.session_state.data_explorer)
            df_final, msg = engineer.build_full_dataset()
            if df_final is not None:
                st.session_state['ml_features'] = df_final
                return {"status": "succès", "details": {"lignes": len(df_final), "features": len(df_final.columns)}}
            return {"status": "erreur", "message": msg}
        except Exception as e: return {"status": "crash", "message": str(e)}

    def _train_models_wrapper(self, args=None):
        try:
            if 'ml_features' not in st.session_state: return {"status": "erreur", "message": "Features manquantes."}
            forecaster = HospitalForecaster(st.session_state['ml_features'])
            forecaster.prepare_data(target_col='nb_admissions')
            metrics = forecaster.train_all_models()
            st.session_state['forecaster'] = forecaster
            st.session_state['metrics'] = metrics
            best = forecaster.best_model_name
            
            # On retourne les détails du meilleur modèle
            return {"status": "succès", "details": {"best_model": best, "r2": metrics[best]['R²'], "mae": metrics[best]['MAE']}}
        except Exception as e: return {"status": "crash", "message": str(e)}

    def _get_forecast_wrapper(self, args=None):
        try:
            if 'forecaster' not in st.session_state: return {"status": "erreur", "message": "IA non entraînée."}
            forecaster = st.session_state['forecaster']
            dates, values, conf = forecaster.generate_future_forecast_recursive(months=3)
            alerts = forecaster.detect_future_anomalies(values)
            trend_pct = ((values[-1] - values[0]) / values[0]) * 100
            
            return {
                "status": "succès", 
                "details": {
                    "tendance_pct": trend_pct,
                    "max_val": float(np.max(values)),
                    "alertes": alerts
                }
            }
        except Exception as e: return {"status": "crash", "message": str(e)}

    def _get_kpis_wrapper(self, args=None):
        try:
            if 'ml_features' not in st.session_state or 'forecaster' not in st.session_state:
                return {"status": "erreur", "message": "Le Cockpit nécessite les données ML et le modèle entraîné."}
            
            monitor = RealtimeMonitor(st.session_state['ml_features'], st.session_state['forecaster'])
            state = monitor.current_state 
            
            # Extraction propre du service en tension
            services_breakdown = monitor.get_service_breakdown()
            if services_breakdown is not None and not services_breakdown.empty:
                # On s'assure que c'est une string propre
                top_tension_service = str(services_breakdown.index[0])
            else:
                top_tension_service = "Global"
            
            state['top_tension_service'] = top_tension_service
            
            return {"status": "succès", "details": state}
        except Exception as e: return {"status": "crash", "message": str(e)}

    def _get_data_exploration_wrapper(self, args=None):
        try:
            analyzer = HospitalAdvancedAnalyzer(self.gspread_client, "bdd_mvp_faker")
            if not analyzer.load_data(): return {"status": "erreur", "message": "Données historiques inaccessibles."}
            
            df_pat = analyzer.get_df('patients')
            df_sej = analyzer.get_df('sejours')
            stats = {}
            
            if df_pat is not None:
                df_pat['age'] = pd.to_numeric(df_pat['age'], errors='coerce')
                stats['demographie'] = {
                    'total': len(df_pat),
                    'age_moyen': round(df_pat['age'].mean(), 1),
                    'senior_ratio': round((df_pat['age'] > 65).mean() * 100, 1),
                    'sex_ratio': df_pat['sexe'].value_counts(normalize=True).to_dict()
                }
            
            if df_sej is not None:
                stats['pathologie'] = df_sej['pathologie_code'].mode()[0] if not df_sej.empty else "N/A"
                
            return {"status": "succès", "details": stats}
        except Exception as e: return {"status": "crash", "message": str(e)}

    def _get_dq_stats_wrapper(self, args=None):
        try:
            raw_data = load_from_google_sheets(self.gspread_client, "bdd_mvp_faker")
            dq = HospitalDataQualityAnalyzer(raw_data)
            scores = dq.compute_global_score()
            details_anomalies = {k: int(v['anomalies']) for k, v in dq.metrics.items() if 'anomalies' in v}
            
            return {
                "status": "succès", 
                "details": {
                    "scores": scores,
                    "anomalies_par_table": details_anomalies
                }
            }
        except Exception as e: return {"status": "crash", "message": str(e)}

    # --- CERVEAU ANALYTIQUE (Le Cœur du Système) ---
    def _generate_human_response(self, tool_name, d):
        """
        Traduit les données techniques en analyse de direction.
        Structure: Constat -> Analyse -> Stratégie.
        """
        
        # 1. MONITORING TEMPS RÉEL (COCKPIT)
        if tool_name == "get_hospital_kpis":
            taux = d.get('taux_occupation', 0) * 100
            trend = d.get('trend_occupation', 0)
            score = min(d.get('health_score', 100), 100) # Clamping pour éviter 182/100
            tension_rh = d.get('tension_personnel', 0)
            service_top = d.get('top_tension_service', 'N/A')
            
            msg = "### 🚁 Rapport de Situation Temps Réel\n\n"
            msg += f"**État Global :** Score de santé de **{score:.0f}/100**.\n"
            msg += f"- Occupation : **{taux:.1f}%** (Tendance 72h: {trend:+.1f}%)\n"
            msg += f"- Tension RH : **{tension_rh:.2f}** (Seuil critique > 1.2)\n\n"
            
            msg += "**🧠 Analyse :** "
            if taux > 90: msg += "Saturation imminente. L'inertie des admissions risque de bloquer les urgences. "
            elif taux > 80: msg += "Activité dense. Les capacités tampons sont entamées. "
            else: msg += "Situation fluide, capacités d'absorption disponibles. "
            
            msg += f"Le point de vigilance est le service **{service_top}**.\n\n"
            
            msg += "**🎯 Recommandations :**\n"
            if score < 70:
                msg += f"1. Priorité absolue : Désengorger le service {service_top}.\n"
                msg += "2. Activer les lits supplémentaires (Plan Niveau 1).\n"
            else:
                msg += "1. Maintenir la surveillance standard.\n2. Optimiser les sorties pour lisser la charge de demain."
            return msg

        # 2. PRÉDICTIONS (FORECASTING)
        elif tool_name == "get_future_forecast":
            tendance = d['tendance_pct']
            max_val = d['max_val']
            alertes = d['alertes']
            
            msg = "### 🔮 Note Prospective (3 Mois)\n\n"
            direction = "HAUSSE" if tendance > 0 else "BAISSE"
            msg += f"**Trajectoire :** Anticipation d'une **{direction} de {abs(tendance):.1f}%**.\n"
            msg += f"**Pic attendu :** {max_val:.0f} admissions/jour.\n\n"
            
            msg += "**🧠 Analyse des Risques :** "
            if len(alertes) > 0:
                msg += f"Je détecte **{len(alertes)} anomalies statistiques** à venir. "
                msg += f"La première tension critique arrivera à **J+{alertes[0]['day']}**.\n\n"
            else:
                msg += "Aucune anomalie majeure détectée. Le flux reste dans le tunnel de normalité.\n\n"
            
            msg += "**🎯 Actions :**\n"
            if tendance > 10:
                msg += "1. Valider les plannings RH pour couvrir la hausse.\n2. Augmenter les stocks de consommables (+15%)."
            else:
                msg += "1. Optimiser le taux de remplissage des lits programmés."
            return msg

        # 3. QUALITÉ DES DONNÉES (AUDIT)
        elif tool_name == "get_data_quality_report":
            scores = d['scores']
            anomalies_dict = d['anomalies_par_table']
            total_anomalies = sum(anomalies_dict.values()) # Correction du bug "not defined"
            
            msg = "### 🛡️ Audit Qualité des Données\n\n"
            msg += f"**Indice de Confiance : {scores['global']}/100**\n"
            msg += f"- Complétude : {scores['completeness']:.1f}%\n- Exactitude : {scores['accuracy']:.1f}%\n\n"
            
            msg += "**🧠 Diagnostic :** "
            if total_anomalies > 0:
                worst_table = max(anomalies_dict, key=anomalies_dict.get)
                msg += f"J'ai identifié **{total_anomalies} anomalies** qui 'bruitent' les analyses. "
                msg += f"La source principale d'erreur est la table **'{worst_table}'** ({anomalies_dict[worst_table]} erreurs).\n\n"
            else:
                msg += "Les données sont propres et fiables pour l'IA.\n\n"
            
            msg += "**🎯 Remédiation :**\n"
            if scores['global'] < 85:
                msg += "1. Demander une purge des doublons et erreurs de dates.\n2. Relancer l'ETL après correction."
            else:
                msg += "Rien à signaler. Maintenir les processus actuels."
            return msg

        # 4. EXPLORATION (HISTORIQUE) - NOUVEAU BLOC AJOUTÉ
        elif tool_name == "analyze_patient_stats":
            demo = d.get('demographie', {})
            patho = d.get('pathologie', 'N/A')
            
            msg = "### 🔬 Analyse Démographique & Clinique\n\n"
            msg += f"**Cohorte :** {demo.get('total', 0)} patients analysés.\n"
            msg += f"**Profil Type :** Âge moyen de **{demo.get('age_moyen', 0)} ans**.\n"
            msg += f"**Gériatrie :** {demo.get('senior_ratio', 0)}% des patients ont plus de 65 ans.\n"
            msg += f"**Pathologie Dominante :** {patho}.\n\n"
            
            msg += "**🧠 Interprétation :** "
            if demo.get('senior_ratio', 0) > 40:
                msg += "La population est vieillissante, ce qui implique des durées de séjour plus longues et une consommation de soins plus élevée. "
            else:
                msg += "La population est relativement jeune, favorisant un turn-over rapide (ambulatoire). "
            
            msg += "\n\n**🎯 Impact Stratégique :**\n"
            msg += "1. Adapter les ratios soignants en fonction de la lourdeur des soins (liée à l'âge).\n"
            msg += "2. Vérifier l'adéquation des équipements pour la pathologie dominante."
            return msg

        # 5. ETL (Mise à jour) - NOUVEAU BLOC
        elif tool_name == "run_etl_pipeline":
            return f"✅ **ETL Terminé :** Base de données synchronisée.\n- Lignes traitées : {d['total_rows_clean']:,}\n- Qualité : Nettoyage effectué avec succès."

        # 6. TRAINING (IA) - NOUVEAU BLOC
        elif tool_name == "train_forecasting_models":
            r2 = d['r2'] * 100
            return f"✅ **IA Mise à Jour :** Le modèle champion est **{d['best_model']}**.\n- Fiabilité (R²) : **{r2:.1f}%**\n- Précision (MAE) : +/- {d['mae']:.1f} admissions/jour."

        # 7. FEATURES - NOUVEAU BLOC
        elif tool_name == "run_feature_engineering":
            return f"✅ **Préparation ML :** Dataset généré.\n- {d['lignes']:,} vecteurs d'apprentissage\n- {d['features']} variables explicatives créées."

        return "Commande exécutée avec succès."

    # --- CERVEAU PRINCIPAL ---
    def ask(self, user_message, chat_history=[]):
        system_prompt = f"""
        ROLE: Tu es Jasmine, la Directrice des Opérations (COO) Virtuelle de l'hôpital.
        TON: Autoritaire mais bienveillant, précis, orienté solution.
        
        OUTILS DISPONIBLES:
        {chr(10).join([f"- {k}: {v['desc']}" for k, v in self.tools_map.items()])}
        
        RÈGLE D'OR:
        Chaque réponse doit contenir :
        1. Les Faits (Chiffres)
        2. L'Analyse (Pourquoi ?)
        3. La Stratégie (Quoi faire ?)
        
        FORMAT DE RÉPONSE JSON:
        {{
          "thought": "Raisonnement interne...",
          "tool": "nom_outil",
          "args": {{}},
          "answer": "Texte si pas d'outil"
        }}
        """

        response_payload = None
        used_model = None

        for model_name in self.MODEL_CASCADE:
            try:
                if "gemini" in model_name:
                    if not self.google_key: continue
                    model = genai.GenerativeModel(model_name)
                    res = model.generate_content(f"SYSTEM: {system_prompt}\nUSER: {user_message}", generation_config={"response_mime_type": "application/json"})
                    response_payload = res.text
                else:
                    if not self.groq_client: continue
                    completion = self.groq_client.chat.completions.create(model=model_name, messages=[{"role":"system","content":system_prompt},{"role":"user","content":user_message}], response_format={"type":"json_object"})
                    response_payload = completion.choices[0].message.content
                used_model = model_name
                break 
            except: continue

        if not response_payload: return "❌ **Erreur Critique :** Services cognitifs indisponibles.", None, None

        try:
            decision = json.loads(re.sub(r'```json\s*|\s*```', '', response_payload).strip())
            tool_name = decision.get("tool")
            thought = f"🧠 Modèle: {used_model}\n💭 {decision.get('thought', '')}"
            
            if tool_name in self.tools_map:
                res = self.tools_map[tool_name]["func"]({})
                if res["status"] == "succès":
                    # Appel de la fonction de réponse humaine enrichie
                    final_response = self._generate_human_response(tool_name, res["details"])
                else:
                    final_response = f"⚠️ **Alerte Technique :** {res.get('message')}"
                return final_response, thought, res
            
            return decision.get("answer", "Je reste en veille."), thought, None
        except Exception as e: return f"Erreur système interne : {str(e)}", None, None

def show_terminal(client, sheet_name="bdd_mvp_faker"):
    apply_custom_styles()
    st.markdown("## 🤖 AIP COMMAND TERMINAL")
    if 'agent_instance' not in st.session_state: st.session_state.agent_instance = JasmineAgent(client)
    if "chat_history" not in st.session_state: st.session_state.chat_history = []

    for msg in st.session_state.chat_history:
        role_class = "user-msg" if msg["role"] == "user" else "bot-msg"
        st.markdown(f'<div class="{role_class}">{msg["content"]}</div>', unsafe_allow_html=True)
        if msg.get("thought"):
            with st.expander("👁️ AIP LOGIC TRACE"): st.markdown(f'<div class="debug-block">{msg["thought"]}</div>', unsafe_allow_html=True)
        if msg.get("data"):
            with st.expander("💾 SYSTEM PAYLOAD"): st.json(msg["data"])
    
    user_input = st.chat_input("Ordre à l'IA (ex: 'Audit Qualité', 'Situation RH', 'Analyse Patients')...")
    if user_input:
        st.session_state.chat_history.append({"role": "user", "content": user_input})
        st.rerun()

    if st.session_state.chat_history and st.session_state.chat_history[-1]["role"] == "user":
        with st.spinner("Analyse stratégique en cours..."):
            response, thought, data = st.session_state.agent_instance.ask(st.session_state.chat_history[-1]["content"])
            st.session_state.chat_history.append({"role": "assistant", "content": response, "thought": thought, "data": data})
            st.rerun()