import os
import pyttsx3
import base64
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import io
import traceback
import re

from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langchain_core.globals import set_llm_cache
from langchain_community.cache import InMemoryCache
from .security import SecurityManager

# Conditional Imports
try: from langchain_ollama import ChatOllama
except ImportError: ChatOllama = None
try: from langchain_openai import ChatOpenAI
except ImportError: ChatOpenAI = None
try: from langchain_anthropic import ChatAnthropic
except ImportError: ChatAnthropic = None
try: from langchain_google_genai import ChatGoogleGenerativeAI
except ImportError: ChatGoogleGenerativeAI = None
try: from langchain_groq import ChatGroq
except ImportError: ChatGroq = None
try: from langchain_mistralai import ChatMistralAI
except ImportError: ChatMistralAI = None

class LLMEngine:
    def __init__(self, 
                 model_name="llama3.2", 
                 provider="auto", 
                 api_key=None, 
                 base_url=None, 
                 temperature=0.1,
                 use_cache=False,
                 security_filters=None,
                 mcp_tools=None,
                 **llm_kwargs):
        
        self.model_name = model_name
        self.provider = provider.lower()
        self.temperature = temperature
        self.api_key = api_key
        self.base_url = base_url
        self.llm_kwargs = llm_kwargs
        self.mcp_tools = mcp_tools or []
        
        self.security = SecurityManager(filters=security_filters)
        
        if use_cache:
            set_llm_cache(InMemoryCache())
            print("        LLM cache enabled.")
        else:
            set_llm_cache(None)

        try: 
            self.tts_engine = pyttsx3.init()
            self.tts_engine.setProperty('rate', 170)
        except: self.tts_engine = None

        if self.provider == "auto":
            self.provider = self._detect_provider(model_name)
        
        print(f"         Engine enabled: {self.provider.upper()} ({self.model_name}) | Security: {list(self.security.active_patterns.keys())}")
        self.llm = self._create_llm_instance()

    def _detect_provider(self, name):
        name = name.lower()
        if "gpt" in name or "o1-" in name: return "openai"
        if "claude" in name: return "anthropic"
        if "gemini" in name: return "google"
        if "mixtral" in name or "llama3-70b" in name: return "groq"
        if "mistral" in name: return "mistral"
        if "deepseek" in name: return "deepseek"
        if "grok" in name: return "grok"
        return "ollama"

    def _create_llm_instance(self):
        kwargs = self.llm_kwargs
        if self.provider == "ollama":
            if not ChatOllama: raise ImportError("pip install langchain-ollama")
            return ChatOllama(model=self.model_name, temperature=self.temperature, base_url=self.base_url or "http://localhost:11434", **kwargs)
        elif self.provider == "openai":
            if not ChatOpenAI: raise ImportError("pip install langchain-openai")
            return ChatOpenAI(model=self.model_name, api_key=self.api_key or os.getenv("OPENAI_API_KEY"), temperature=self.temperature, **kwargs)
        elif self.provider == "anthropic":
            if not ChatAnthropic: raise ImportError("pip install langchain-anthropic")
            return ChatAnthropic(model_name=self.model_name, api_key=self.api_key or os.getenv("ANTHROPIC_API_KEY"), temperature=self.temperature, **kwargs)
        elif self.provider == "google":
            if not ChatGoogleGenerativeAI: raise ImportError("pip install langchain-google-genai")
            return ChatGoogleGenerativeAI(model=self.model_name, google_api_key=self.api_key or os.getenv("GOOGLE_API_KEY"), temperature=self.temperature, **kwargs)
        elif self.provider == "groq":
            if not ChatGroq: raise ImportError("pip install langchain-groq")
            return ChatGroq(model_name=self.model_name, api_key=self.api_key or os.getenv("GROQ_API_KEY"), temperature=self.temperature, **kwargs)
        elif self.provider == "mistral":
            if not ChatMistralAI: raise ImportError("pip install langchain-mistralai")
            return ChatMistralAI(model=self.model_name, api_key=self.api_key or os.getenv("MISTRAL_API_KEY"), temperature=self.temperature, **kwargs)
        elif self.provider in ["deepseek", "grok", "custom", "xai"]:
            if not ChatOpenAI: raise ImportError("pip install langchain-openai")
            url = self.base_url
            key = self.api_key
            if self.provider == "deepseek" and not url: url, key = "https://api.deepseek.com/v1", key or os.getenv("DEEPSEEK_API_KEY")
            elif self.provider == "grok" and not url: url, key = "https://api.x.ai/v1", key or os.getenv("XAI_API_KEY")
            return ChatOpenAI(model=self.model_name, api_key=key, base_url=url, temperature=self.temperature, **kwargs)
        else: raise ValueError(f"Unknown provider: {self.provider}")

    def _build_system_prompt(self, config):
        parts = []
        parts.append(f"ROLE: {config.get('role', 'AI Assistant')}")
        parts.append(f"GOAL: {config.get('goal', 'Help the user')}")
        if config.get("company_name"):
            c_info = f"COMPANY: {config.get('company_name')}"
            if config.get("company_description"): c_info += f" - {config.get('company_description')}"
            if config.get("company_url"): c_info += f" ({config.get('company_url')})"
            parts.append(c_info)
        if config.get("user_profile"):
            parts.append(f"USER PROFILE: {config.get('user_profile')}")
            parts.append("IMPORTANT: Enforce access control based on User Profile.")
        if config.get("instructions"): parts.append(f"INSTRUCTIONS: {config.get('instructions')}")
        parts.append("LANGUAGE RULE: Respond in the SAME LANGUAGE as the user.")
        if config.get("reflection"): parts.append("MODE: Reflection Enabled. Think step-by-step.")
        if config.get("system_prompt"): parts.append(f"SYSTEM OVERRIDE: {config.get('system_prompt')}")
        return "\n".join(parts)

    def generate(self, prompt, context=None, context_config={}, image_path=None, vocal_out=False, stream=False, output_format="text"):
        messages = []
        sys_prompt = self._build_system_prompt(context_config)
        messages.append(SystemMessage(content=sys_prompt))

        format_instruction = ""
        if output_format == "json": format_instruction = "\nIMPORTANT: Respond only in valid JSON."
        elif output_format == "markdown": format_instruction = "\nIMPORTANT: Use Markdown format."
        
        if self.security.is_active:
            filters_list = ", ".join(self.security.active_patterns.keys())
            format_instruction += f"\nSECURITY: Mask these types: [{filters_list}]."

        final_prompt = prompt + format_instruction
        if context:
            final_prompt = (f"Use the context below.\n--- CONTEXT ---\n{context}\n----------------\nQUESTION: {final_prompt}")
            
        content = [{"type": "text", "text": final_prompt}]

        if image_path:
            try:
                with open(image_path, "rb") as f:
                    enc = base64.b64encode(f.read()).decode("utf-8")
                content.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{enc}"}})
            except: pass
            
        messages.append(HumanMessage(content=content))
        
        if stream: return self._stream_response_generator(messages, vocal_out)
        else:
            try:
                response_obj = self.llm.invoke(messages)
                clean_text = self.security.scrub(response_obj.content, context_config.get('user_profile'))
                if vocal_out and self.tts_engine: self.speak(clean_text)
                return clean_text
            except Exception as e: return f"  Error : {e}"

    # =========================================================================
    # DATA ANALYST AGENT (SELF-HEALING & ROBUST)
    # =========================================================================
    def generate_analyst(self, query, df, context_config, history="", output_format="text", stream=False, canvas=None):
        """
        Executes Python code with SELF-CORRECTION loop.
        If 'could not convert string to float' occurs, it feeds the error back to LLM to fix it.
        """
        buffer = io.StringIO()
        df.info(buf=buffer)
        df_info = buffer.getvalue()
        df_head = df.head(3).to_markdown()
        
        sys_prompt = self._build_system_prompt(context_config)
        
        # PROMPT RENFORCÉ POUR PRÉVENIR LES ERREURS
        code_prompt = f"""
        You are a Senior Data Scientist. You have a pandas DataFrame `df`.
        
        DATAFRAME INFO:
        {df_info}
        
        FIRST 3 ROWS:
        {df_head}
        
        USER REQUEST: {query}
        HISTORY: {history}
        
        MANDATORY RULES FOR ROBUST CODE:
        1. **PREPROCESSING IS REQUIRED**: Before doing correlation (.corr), regression, or math, YOU MUST DROP non-numeric columns OR encode them (pd.get_dummies).
           - Example: `df_numeric = df.select_dtypes(include=[np.number])`
           - Example: `df_encoded = pd.get_dummies(df, drop_first=True)`
        2. **HANDLE NaNs**: Use `df.dropna()` or `df.fillna(0)` before plotting.
        3. **CHARTS**: Use `matplotlib` or `seaborn`. Save to 'latest_chart.png'.
        4. **OUTPUT**: Print the final result.
        
        Respond ONLY with the Python code block.
        """
        
        messages = [
            SystemMessage(content=sys_prompt),
            HumanMessage(content=code_prompt)
        ]
        
        # BOUCLE DE CORRECTION (MAX 3 TENTATIVES)
        max_retries = 3
        last_error = None
        code_to_run = ""
        execution_result = ""
        
        for attempt in range(max_retries):
            try:
                print(f"    [Analyst] Thinking (Attempt {attempt+1}/{max_retries})...")
                resp = self.llm.invoke(messages)
                raw_response = resp.content
                
                # Extraction Code
                code_match = re.search(r"```python(.*?)```", raw_response, re.DOTALL)
                code_to_run = code_match.group(1).strip() if code_match else raw_response
                
                # Exécution
                import sys
                old_stdout = sys.stdout
                redirected_output = io.StringIO()
                sys.stdout = redirected_output
                
                # Environnement d'exécution
                local_scope = {"df": df, "pd": pd, "np": np, "plt": plt, "sns": sns}
                
                exec(code_to_run, {}, local_scope)
                execution_result = redirected_output.getvalue()
                sys.stdout = old_stdout
                
                # Si on arrive ici, le code a fonctionné !
                break
            
            except Exception as e:
                # Capture de
                sys.stdout = old_stdout # Restauration console importante
                last_error = str(e)
                print(f"    [Analyst] Error detected: {last_error}. Retrying...")
                
                # On injecte l'erreur dans la conversation pour que le LLM se corrige
                error_feedback = f"""
                EXECUTION ERROR: {last_error}
                
                The previous code failed. 
                HINT: If the error is "could not convert string to float", it means you tried to calculate correlation or math on Text/String columns.
                FIX: Select only numeric columns: `df_numeric = df.select_dtypes(include=[np.number])` before calculating.
                
                Please rewrite the code to fix this error.
                """
                messages.append(AIMessage(content=code_to_run))
                messages.append(HumanMessage(content=error_feedback))
        
        # Fin de la boucle
        if not execution_result and last_error:
            return f"I tried to analyze the data 3 times but failed. Last error: {last_error}"

        # 5. Interprétation finale du résultat
        interpretation_prompt = f"""
        USER REQUEST: {query}
        
        EXECUTED CODE:
        {code_to_run}
        
        EXECUTION RESULT:
        {execution_result}
        
        Based on the execution result, answer the user's question naturally in their language.
        If a chart was generated (latest_chart.png), mention "I have generated the chart."
        Do not show the code, just the analysis.
        """
        
        final_messages = [
            SystemMessage(content=sys_prompt),
            HumanMessage(content=interpretation_prompt)
        ]
        
        if stream:
            return self._stream_response_generator(final_messages, vocal_out=False)
        else:
            final_resp = self.llm.invoke(final_messages)
            clean_text = self.security.scrub(final_resp.content, context_config.get('user_profile'))
            
            # Canvas Integration (Simulation)
            if canvas and "latest_chart.png" in execution_result and os.path.exists("latest_chart.png"):
                print("    [Canvas] Updating visual interface with chart...")
            
            return clean_text

    def _stream_response_generator(self, messages, vocal_out):
        """Generator with BUFFER to secure streaming."""
        full_text_buffer = ""
        display_buffer = ""
        
        try:
            for chunk in self.llm.stream(messages):
                token = chunk.content
                if token:
                    display_buffer += token
                    if len(display_buffer) > 25 or token in [".", "\n", "!", "?", " "]:
                        clean_chunk = self.security.scrub(display_buffer)
                        yield clean_chunk
                        full_text_buffer += clean_chunk
                        display_buffer = ""
            
            if display_buffer:
                clean_chunk = self.security.scrub(display_buffer)
                full_text_buffer += clean_chunk
                yield clean_chunk

            if vocal_out and self.tts_engine: self.speak(full_text_buffer)
            
        except Exception as e:
            yield f"  Stream Error: {e}"

    def speak(self, text):
        try:
            self.tts_engine.say(text)
            self.tts_engine.runAndWait()
        except: pass
            
# # ------------------------------------------------
# import os
# import pyttsx3
# from langchain_core.messages import HumanMessage
# # from langchain.globals import set_llm_cache
# from langchain_core.globals import set_llm_cache
# from langchain_community.cache import InMemoryCache
# from .security import SecurityManager

# # Imports conditionnels
# try: from langchain_ollama import ChatOllama
# except ImportError: ChatOllama = None
# try: from langchain_openai import ChatOpenAI
# except ImportError: ChatOpenAI = None
# try: from langchain_anthropic import ChatAnthropic
# except ImportError: ChatAnthropic = None
# try: from langchain_google_genai import ChatGoogleGenerativeAI
# except ImportError: ChatGoogleGenerativeAI = None
# try: from langchain_groq import ChatGroq
# except ImportError: ChatGroq = None
# try: from langchain_mistralai import ChatMistralAI
# except ImportError: ChatMistralAI = None

# class LLMEngine:
#     def __init__(self, 
#                  model_name="llama3.2", 
#                  provider="auto", 
#                  api_key=None, 
#                  base_url=None, 
#                  temperature=0.1,
#                  use_cache=False,
#                  security_filters=None, # <--- Remplace security_active (peut être True ou liste)
#                  **llm_kwargs):
        
#         self.model_name = model_name
#         self.provider = provider.lower()
#         self.temperature = temperature
#         self.api_key = api_key
#         self.base_url = base_url
#         self.llm_kwargs = llm_kwargs
        
#         # Init Sécurité avec les filtres spécifiques
#         self.security = SecurityManager(filters=security_filters)
        
#         # Init Cache
#         if use_cache:
#             set_llm_cache(InMemoryCache())
#             print("🚀 LLM cache enabled.")
#         else:
#             set_llm_cache(None)

#         try:
#             self.tts_engine = pyttsx3.init()
#             self.tts_engine.setProperty('rate', 170)
#         except: self.tts_engine = None

#         if self.provider == "auto":
#             self.provider = self._detect_provider(model_name)
        
#         print(f"🤖 Engine enabled: {self.provider.upper()} ({self.model_name}) | Security: {list(self.security.active_patterns.keys()) if self.security.is_active else 'OFF'}")
#         self.llm = self._create_llm_instance()

#     def _detect_provider(self, name):
#         name = name.lower()
#         if "gpt" in name or "o1-" in name: return "openai"
#         if "claude" in name: return "anthropic"
#         if "gemini" in name: return "google"
#         if "mixtral" in name or "llama3-70b" in name: return "groq"
#         if "mistral" in name: return "mistral"
#         if "deepseek" in name: return "deepseek"
#         if "grok" in name: return "grok"
#         return "ollama"

#     def _create_llm_instance(self):
#         if self.provider == "ollama":
#             if not ChatOllama: raise ImportError("pip install langchain-ollama")
#             return ChatOllama(model=self.model_name, temperature=self.temperature, base_url=self.base_url or "http://localhost:11434", **self.llm_kwargs)
#         elif self.provider == "openai":
#             if not ChatOpenAI: raise ImportError("pip install langchain-openai")
#             return ChatOpenAI(model=self.model_name, api_key=self.api_key or os.getenv("OPENAI_API_KEY"), temperature=self.temperature, **self.llm_kwargs)
        # elif self.provider == "anthropic":
        #     if not ChatAnthropic: raise ImportError("pip install langchain-anthropic")
        #     return ChatAnthropic(model_name=self.model_name, api_key=self.api_key or os.getenv("ANTHROPIC_API_KEY"), temperature=self.temperature, **self.llm_kwargs)
        # elif self.provider == "google":
        #     if not ChatGoogleGenerativeAI: raise ImportError("pip install langchain-google-genai")
        #     return ChatGoogleGenerativeAI(model=self.model_name, google_api_key=self.api_key or os.getenv("GOOGLE_API_KEY"), temperature=self.temperature, **self.llm_kwargs)
        # elif self.provider == "groq":
        #     if not ChatGroq: raise ImportError("pip install langchain-groq")
        #     return ChatGroq(model_name=self.model_name, api_key=self.api_key or os.getenv("GROQ_API_KEY"), temperature=self.temperature, **self.llm_kwargs)
        # elif self.provider == "mistral":
        #     if not ChatMistralAI: raise ImportError("pip install langchain-mistralai")
        #     return ChatMistralAI(model=self.model_name, api_key=self.api_key or os.getenv("MISTRAL_API_KEY"), temperature=self.temperature, **self.llm_kwargs)
        # elif self.provider in ["deepseek", "grok", "custom", "xai"]:
        #     if not ChatOpenAI: raise ImportError("pip install langchain-openai")
        #     url = self.base_url
        #     key = self.api_key
        #     if self.provider == "deepseek" and not url: url, key = "https://api.deepseek.com/v1", key or os.getenv("DEEPSEEK_API_KEY")
        #     elif self.provider == "grok" and not url: url, key = "https://api.x.ai/v1", key or os.getenv("XAI_API_KEY")
        #     return ChatOpenAI(model=self.model_name, api_key=key, base_url=url, temperature=self.temperature, **self.llm_kwargs)
        # else: raise ValueError(f"Provider inconnu : {self.provider}")

#     def generate(self, prompt, context=None, image_path=None, vocal_out=False, stream=False, output_format="text"):
#         messages = []
        
#         format_instruction = ""
#         if output_format == "json": format_instruction = "\nIMPORTANT: Respond only in valid JSON."
#         elif output_format == "markdown": format_instruction = "\nIMPORTANT: Use Markdown format."
#         elif output_format == "toon": format_instruction = "\nSTYLE: You are a cartoon character!"
        
#         # Instruction système de sécurité dynamique
#         if self.security.is_active:
#             filters_list = ", ".join(self.security.active_patterns.keys())
#             format_instruction += f"\nSECURITY: The following data is CONFIDENTIAL: [{filters_list}]. Never display them. Replace them with a note indicating they are protected."

#         final_prompt = prompt + format_instruction
#         if context:
#             final_prompt = (
#                 "You are an expert assistant. Use the context below.\n"
#                 f"--- CONTEXT ---\n{context}\n----------------\n"
#                 f"QUESTION: {final_prompt}"
#             )
            
#         content = [{"type": "text", "text": final_prompt}]

#         if image_path:
#             import base64
#             try:
#                 with open(image_path, "rb") as f:
#                     enc = base64.b64encode(f.read()).decode("utf-8")
#                 content.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{enc}"}})
#             except: pass
            
#         messages.append(HumanMessage(content=content))
        
#         if stream:
#             return self._stream_response_generator(messages, vocal_out)
#         else:
#             try:
#                 response_obj = self.llm.invoke(messages)
#                 clean_text = self.security.scrub(response_obj.content)
#                 if vocal_out and self.tts_engine: self.speak(clean_text)
#                 return clean_text
#             except Exception as e:
#                 return f"❌ Erreur : {e}"

#     def _stream_response_generator(self, messages, vocal_out):
#         """Generator with BUFFER to secure streaming."""
#         full_text_buffer = ""
#         display_buffer = ""
        
#         try:
#             for chunk in self.llm.stream(messages):
#                 token = chunk.content
#                 if token:
#                     display_buffer += token
                    
#                     # On affiche quand on a assez de contexte pour que les regex fonctionnent
#                     if len(display_buffer) > 25 or token in [".", "\n", "!", "?", " "]:
#                         clean_chunk = self.security.scrub(display_buffer)
#                         yield clean_chunk
#                         full_text_buffer += clean_chunk
#                         display_buffer = ""
            
#             # Fin du buffer
#             if display_buffer:
#                 clean_chunk = self.security.scrub(display_buffer)
#                 full_text_buffer += clean_chunk
#                 yield clean_chunk

#             if vocal_out and self.tts_engine: self.speak(full_text_buffer)
            
#         except Exception as e:
#             yield f"❌ Stream Error: {e}"

#     def speak(self, text):
#         try:
#             self.tts_engine.say(text)
#             self.tts_engine.runAndWait()
#         except: pass