import os
import threading
import pandas as pd
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

# Internal Imports
from .llm_engine import LLMEngine
from .vector_store import VectorDBFactory
from .loaders import UniversalLoader
from .watcher import RostaingWatcher, PollingWatcher
from .embeddings import EmbeddingsManager

class RostaingBrain:
    def __init__(self, 
                 # --- CORE PARAMETERS ---
                 llm_model="llama3.2", 
                 llm_provider="auto",
                 llm_api_key=None,
                 llm_base_url=None,
                 
                 embedding_model="BAAI/bge-small-en-v1.5",
                 embedding_source="fastembed",
                 vector_db="chroma",   
                 data_source="./data", 
                 
                 # --- AUTOMATION ---
                 auto_update=True,
                 poll_interval=60,
                 reset_db=False,
                 memory=False,

                 # --- LLM CONFIG (MOVED HERE) ---
                 temperature=0.1,
                 max_tokens=None,
                 top_p=None,
                 top_k=None,
                 seed=None,          
                 stream=False,       
                 cache=True,
                 output_format="text", # Global output format
                 
                 # --- AGENT IDENTITY & CONTEXT ---
                 role="Helpful AI Assistant",
                 goal="Assist the user accurately.",
                 instructions="Answer concisely and exhaustively.",
                 reflection=False,
                 
                 # --- COMPANY CONTEXT ---
                 company_name=None,
                 company_description=None,
                 company_url=None,

                 # --- SECURITY & MULTI-USER ---
                 security_filters=None, 
                 user_profile=None,
                 user_id=None,
                 session_id=None,
                 agent_id=None,
                 system_prompt=None,
                 
                 # --- TOOLS & UI ---
                 mcp_tools=None,
                 canvas=None
                 ):
        
        # Configuration Store
        self.config = {
            "role": role,
            "goal": goal,
            "instructions": instructions,
            "company_name": company_name,
            "company_description": company_description,
            "company_url": company_url,
            "user_profile": user_profile,
            "user_id": user_id,
            "session_id": session_id,
            "agent_id": agent_id,
            "system_prompt": system_prompt,
            "reflection": reflection,
            "stream": stream,
            "seed": seed,
            "output_format": output_format
        }

        self.db_type = vector_db
        self.memory_object = None
        self.db_config = None
        self.web_url = None
        self.mcp_tools = mcp_tools or []
        self.canvas = canvas
        
        self.use_memory = memory
        self.chat_history = [] 
        
        # Context Flags
        self.data_analyst_mode = False 
        self.dataframe_context = None
        self.vision_mode = False
        self.active_image_path = None

        # --- INTELLIGENT SOURCE DETECTION ---
        if isinstance(data_source, dict):
            self.mode = "database"; self.db_config = data_source; self.watch_target = None
            print(f"          DATABASE MODE: {data_source.get('type', 'sql').upper()}")
            if data_source.get('type') == 'sql': self.data_analyst_mode = True
            
        elif isinstance(data_source, str):
            if data_source.startswith("http"): 
                self.mode = "web"; self.web_url = data_source; print(f"          WEB MODE")
            else:
                self.abs_path = os.path.abspath(data_source)
                ext = os.path.splitext(self.abs_path)[1].lower()
                
                # Check for Image (Vision Mode)
                if ext in ['.jpg', '.jpeg', '.png', '.bmp', '.webp']:
                    self.mode = "image"
                    self.vision_mode = True
                    self.active_image_path = self.abs_path
                    self.watch_target = os.path.dirname(self.abs_path)
                    print(f"     VISION MODE (Image: {os.path.basename(self.abs_path)})")
                    
                # Check for File
                elif ext: 
                    self.mode = "file"; self.watch_target = os.path.dirname(self.abs_path) or os.getcwd(); self.target_filename = os.path.basename(self.abs_path)
                    print(f"     SINGLE FILE MODE")
                    # Check if file is tabular for Analyst Mode
                    if ext in ['.csv', '.xlsx', '.xls', '.parquet', '.json']:
                        self.data_analyst_mode = True
                else: 
                    self.mode = "directory"; self.watch_target = self.abs_path; print(f"   FOLDER MODE")
        else:
            self.mode = "memory"; self.memory_object = data_source; print(f"      MEMORY MODE"); auto_update = False 
            if isinstance(data_source, (pd.DataFrame, pd.Series)):
                self.data_analyst_mode = True
                self.dataframe_context = data_source

        print(f"    Init RostaingBrain ({llm_model}) | Analyst: {self.data_analyst_mode} | Vision: {self.vision_mode}")

        if reset_db: VectorDBFactory.reset("./db_storage")

        self.embed_manager = EmbeddingsManager(model_name=embedding_model, source=embedding_source)
        self.embedding_function = self.embed_manager.get_function()
        self.loader = UniversalLoader()
        self.db = VectorDBFactory.create(vector_db, self.embedding_function)
        
        # LLM Parameters
        llm_kwargs = {}
        if max_tokens: llm_kwargs['max_tokens'] = max_tokens
        if top_p: llm_kwargs['top_p'] = top_p
        if top_k: llm_kwargs['top_k'] = top_k
        if seed: llm_kwargs['seed'] = seed
        
        self.engine = LLMEngine(
            model_name=llm_model, 
            provider=llm_provider, 
            api_key=llm_api_key,
            base_url=llm_base_url,
            temperature=temperature,
            use_cache=cache,
            security_filters=security_filters,
            mcp_tools=self.mcp_tools,
            **llm_kwargs
        )
        
        self._initial_ingest()
        
        if auto_update and self.mode != "memory" and not self.vision_mode:
            if self.mode in ["database", "web"]:
                self.watcher = PollingWatcher(self._on_polling_event, data_source if self.mode == "database" else self.web_url, interval=poll_interval)
                self.watcher.start()
            elif self.mode in ["directory", "file"]:
                if self.mode == "directory" and not os.path.exists(self.watch_target): os.makedirs(self.watch_target)
                self.watcher = RostaingWatcher(self._on_file_event, self.watch_target)
                t = threading.Thread(target=self.watcher.start); t.daemon = True; t.start()
                print(f"       Active monitoring.")

    def _initial_ingest(self):
        if self.mode == "database": self._process_database(self.db_config)
        elif self.mode == "web": self._process_web(self.web_url)
        elif self.mode == "memory": self._ingest_memory_object()
        elif self.vision_mode: pass # No ingestion needed for single image vision yet
        else: self._ingest_existing_files()

    def _process_database(self, c):
        docs, raw_df = self.loader.process_sql(c, return_raw=True) if c.get("type","sql")=="sql" else (self.loader.process_mongodb(c), None)
        if raw_df is not None: self.dataframe_context = raw_df
        self._update_db(docs)

    def _process_web(self, u): docs=self.loader.process_web(u); self._update_db(docs)
    
    def _ingest_memory_object(self):
        docs = self.loader.process_dataframe_object(self.memory_object); self._update_db(docs)

    def _update_db(self, docs):
        if not docs: return
        splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=300)
        chunks = splitter.split_documents(docs)
        self.db = VectorDBFactory.add_docs(self.db, chunks, self.db_type, self.embedding_function, "./db_storage")
        print(f"   RAG updated ({len(chunks)} chunks).")

    def _on_polling_event(self, c): 
        if self.mode=="database": self._process_database(c)
        else: self._process_web(c)
        
    def _on_file_event(self, p):
        p=os.path.abspath(p)
        if self.mode=="file": 
            if p==self.abs_path: self.process_file(p)
        else: self.process_file(p)
        
    def process_file(self, file_path):
        if not os.path.exists(file_path) or "~" in file_path: return
        try: 
            print(f"    {os.path.basename(file_path)}")
            docs, df = self.loader.load_file(file_path, return_raw=True)
            if df is not None: 
                self.dataframe_context = df
                self.data_analyst_mode = True
            self._update_db(docs)
        except: pass

    def _ingest_existing_files(self):
        if self.mode=="file": 
            if os.path.exists(self.abs_path): self.process_file(self.abs_path)
        else:
            if not os.path.exists(self.watch_target): os.makedirs(self.watch_target); return
            for r,_,fs in os.walk(self.watch_target):
                for f in fs: self.process_file(os.path.join(r,f))

    def _safe_retrieval(self, query):
        docs = []
        try:
            if self.db: 
                retriever = self.db.as_retriever(search_type="mmr", search_kwargs={"k": 6, "fetch_k": 20})
                docs = retriever.invoke(query)
        except Exception as e: print(f"    Retrieval Error: {e}")
        return docs

    def chat(self, user_query):
        """
        Unified Agent Chat Endpoint. No extra params.
        """
        # Load config params
        effective_stream = self.config['stream']
        output_fmt = self.config['output_format']
        
        # 1. RAG Retrieval (Skip if Vision Mode only)
        context_text = ""
        if not self.vision_mode:
            docs = self._safe_retrieval(user_query)
            context_text = "\n\n".join([d.page_content for d in docs]) if docs else ""
            if context_text: print(f"     Context used : {len(docs)} segments.")
        
        # 2. History Context
        history_str = ""
        if self.use_memory and self.chat_history:
            history_str = "HISTORY:\n" + "\n".join([f"User:{q}\nAI:{a}" for q,a in self.chat_history[-5:]])
        
        # 3. Execution Routing
        if self.data_analyst_mode and self.dataframe_context is not None:
            # === DATA ANALYST MODE (AUTO-FIX ENABLED) ===
            print("    > Engaging Data Analyst Agent (Python Execution)...")
            response = self.engine.generate_analyst(
                query=user_query,
                df=self.dataframe_context,
                context_config=self.config,
                history=history_str,
                output_format=output_fmt,
                stream=effective_stream,
                canvas=self.canvas
            )
        else:
            # === STANDARD / VISION MODE ===
            full_context = f"{history_str}\n\nRAG CONTEXT:\n{context_text}"
            response = self.engine.generate(
                prompt=user_query,
                context=full_context,
                context_config=self.config,
                image_path=self.active_image_path, # Passed from init config
                stream=effective_stream,
                output_format=output_fmt
            )

        if self.use_memory and not effective_stream: self.chat_history.append((user_query, response))
        return response

# # ---------------------------------------
# import os
# import threading
# import pandas as pd
# # from langchain.text_splitter import RecursiveCharacterTextSplitter
# from langchain_text_splitters import RecursiveCharacterTextSplitter
# # from langchain.docstore.document import Document
# from langchain_core.documents import Document

# # Imports internes
# from .llm_engine import LLMEngine
# from .vector_store import VectorDBFactory
# from .loaders import UniversalLoader
# from .watcher import RostaingWatcher, PollingWatcher
# from .embeddings import EmbeddingsManager

# class RostaingBrain:
#     def __init__(self, 
#                  llm_model="llama3.2", 
#                  llm_provider="auto",
#                  llm_api_key=None,
#                  llm_base_url=None,
#                  embedding_model="BAAI/bge-small-en-v1.5",
#                  embedding_source="fastembed",
#                  vector_db="chroma",   
#                  data_source="./data", 
#                  auto_update=True,
#                  poll_interval=60,
#                  reset_db=False,
#                  memory=False,
#                  # --- PARAMÈTRES LLM ---
#                  temperature=0.1,
#                  max_tokens=None,
#                  top_p=None,
#                  top_k=None,
#                  cache=True,
#                  # --- SÉCURITÉ (NOUVEAU) ---
#                  security_filters=None): # Liste ["BIC", "EMAIL"] ou True/False
        
#         self.db_type = vector_db
#         self.memory_object = None
#         self.db_config = None
#         self.web_url = None
        
#         self.use_memory = memory
#         self.chat_history = [] 
        
#         # --- DÉTECTION SOURCE ---
#         if isinstance(data_source, dict):
#             self.mode = "database"; self.db_config = data_source; self.watch_target = None
#             print(f"🗄️  DATABASE MODE: {data_source.get('type', 'sql').upper()}")
#         elif isinstance(data_source, str):
#             if data_source.startswith("http"): self.mode = "web"; self.web_url = data_source; print(f"🌐 WEB MODE")
#             else:
#                 self.abs_path = os.path.abspath(data_source)
#                 if os.path.splitext(self.abs_path)[1]: self.mode = "file"; self.watch_target = os.path.dirname(self.abs_path) or os.getcwd(); self.target_filename = os.path.basename(self.abs_path); print(f"📄 SINGLE FILE MODE")
#                 else: self.mode = "directory"; self.watch_target = self.abs_path; print(f"📁 FOLDER MODE")
#         else:
#             self.mode = "memory"; self.memory_object = data_source; print(f"📊 MEMORY MODE"); auto_update = False 

#         print(f"🧠 Init RostaingBrain ({llm_model})")

#         if reset_db: VectorDBFactory.reset("./db_storage")

#         self.embed_manager = EmbeddingsManager(model_name=embedding_model, source=embedding_source)
#         self.embedding_function = self.embed_manager.get_function()
#         self.loader = UniversalLoader()
#         self.db = VectorDBFactory.create(vector_db, self.embedding_function)
        
#         # Transmission des params
#         llm_kwargs = {}
#         if max_tokens: llm_kwargs['max_tokens'] = max_tokens
#         if top_p: llm_kwargs['top_p'] = top_p
#         if top_k: llm_kwargs['top_k'] = top_k
        
#         self.engine = LLMEngine(
#             model_name=llm_model, 
#             provider=llm_provider, 
#             api_key=llm_api_key, 
#             base_url=llm_base_url,
#             temperature=temperature,
#             use_cache=cache,
#             security_filters=security_filters, # <--- Passage des filtres
#             **llm_kwargs
#         )
        
#         self._initial_ingest()
        
#         if auto_update and self.mode != "memory":
#             if self.mode in ["database", "web"]:
#                 self.watcher = PollingWatcher(self._on_polling_event, data_source if self.mode == "database" else self.web_url, interval=poll_interval)
#                 self.watcher.start()
#             elif self.mode in ["directory", "file"]:
#                 if self.mode == "directory" and not os.path.exists(self.watch_target): os.makedirs(self.watch_target)
#                 self.watcher = RostaingWatcher(self._on_file_event, self.watch_target)
#                 t = threading.Thread(target=self.watcher.start); t.daemon = True; t.start()
#                 print(f"👁️  Active monitoring.")

#     # ... (Le reste des méthodes reste identique) ...
#     # Je copie ici les méthodes indispensables pour que le code soit complet
#     def _initial_ingest(self):
#         if self.mode == "database": self._process_database(self.db_config)
#         elif self.mode == "web": self._process_web(self.web_url)
#         elif self.mode == "memory": self._ingest_memory_object()
#         else: self._ingest_existing_files()

#     def _process_database(self, c): docs=self.loader.process_sql(c) if c.get("type","sql")=="sql" else (self.loader.process_mongodb(c) if c.get("type")=="mongodb" else self.loader.process_neo4j(c)); self._update_db(docs)
#     def _process_web(self, u): docs=self.loader.process_web(u); self._update_db(docs)
#     def _ingest_memory_object(self): docs=self.loader.process_dataframe_object(self.memory_object); self._update_db(docs)

#     def _update_db(self, docs):
#         if not docs: return
#         splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=300)
#         chunks = splitter.split_documents(docs)
#         self.db = VectorDBFactory.add_docs(self.db, chunks, self.db_type, self.embedding_function, "./db_storage")
#         print(f"✅ RAG updated ({len(chunks)} chunks).")

#     def _on_polling_event(self, c): 
#         if self.mode=="database": self._process_database(c)
#         else: self._process_web(c)
#     def _on_file_event(self, p):
#         p=os.path.abspath(p)
#         if self.mode=="file": 
#             if p==self.abs_path: self.process_file(p)
#         else: self.process_file(p)
#     def process_file(self, file_path):
#         if not os.path.exists(file_path) or "~" in file_path: return
#         try: 
#             print(f"⚙️ {os.path.basename(file_path)}")
#             self._update_db(self.loader.load_file(file_path))
#         except: pass
#     def _ingest_existing_files(self):
#         if self.mode=="file": 
#             if os.path.exists(self.abs_path): self.process_file(self.abs_path)
#         else:
#             if not os.path.exists(self.watch_target): os.makedirs(self.watch_target); return
#             for r,_,fs in os.walk(self.watch_target): 
#                 for f in fs: self.process_file(os.path.join(r,f))

#     def _safe_retrieval(self, query):
#         docs = []
#         try:
#             if self.db: 
#                 retriever = self.db.as_retriever(search_type="mmr", search_kwargs={"k": 6, "fetch_k": 20})
#                 docs = retriever.invoke(query)
#         except Exception as e: print(f"⚠️ Retrieval Error: {e}")
#         return docs

#     def chat(self, user_query, image_path=None, vocal_response=False, stream=False, output_format="text"):
#         docs = self._safe_retrieval(user_query)
#         context = "\n\n".join([d.page_content for d in docs]) if docs else ""
#         if context: print(f"📘 Context used : {len(docs)} segments.")
#         else: print("⚠️  No RAG context found.")
        
#         if self.use_memory and self.chat_history:
#             history = "HISTORY:\n" + "\n".join([f"User:{q}\nAI:{a}" for q,a in self.chat_history[-5:]])
#             context = f"{history}\n\nRAG CONTEXT:\n{context}"

#         response = self.engine.generate(user_query, context=context, image_path=image_path, vocal_out=vocal_response, stream=stream, output_format=output_format)
#         if self.use_memory and not stream: self.chat_history.append((user_query, response))
#         return response