import os
import threading
import logging
import datetime
import pandas as pd
from typing import Callable
try:
    import polars as pl
except ImportError:
    pl = None

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 RostaingAgent:
    """
    RostaingAgent: The core of the autonomous hybrid RAG framework.
    Capable of self-reasoning, task decomposition, and Big Data handling.
    """
    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 & LOGS ---
                 auto_update=True,
                 poll_interval=1800,
                 reset_db=False,
                 memory=True, # Default to True for ChatGPT-like behavior
                 save_graph=False,
                 save_logs=False,

                 # --- LLM CONFIG ---
                 temperature=0.1,
                 max_tokens=4096, # Increased for complex reasoning
                 top_p=None,
                 top_k=None,
                 seed=None,          
                 stream=False,       
                 cache=True,
                 output_format="text",
                 
                 # --- AGENT IDENTITY & CONTEXT ---
                 role="Expert Autonomous Data & RAG Assistant",
                 goal="Provide highly accurate, structured, and factual answers based strictly on the knowledge base.",
                 instructions="Decompose complex tasks. Never hallucinate. If data is missing, state it clearly.",
                 reflection=True, # Default to True for Agentic behavior
                 
                 # --- COMPANY CONTEXT ---
                 company_name=None,
                 company_description=None,
                 company_url=None,

                 # --- SECURITY & MULTI-USER ---
                 security_filters=None, 
                 user_profile="default", # Dynamically used in security.py
                 user_id=None,
                 session_id=None,
                 agent_id=None,
                 system_prompt=None,
                 
                 # --- TOOLS & UI ---
                 mcp_tools=None,
                 canvas=None
                 ):
        
        self.save_graph = save_graph
        self.save_logs = save_logs
        self._setup_directories_and_logging()
        
        self.log(f"Initializing RostaingAgent (v1.2.0) with model: {llm_model}")

        # 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,
            "save_graph": self.save_graph
        }

        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
        
        # ChatGPT-like Memory (Sliding Window Context)
        self.use_memory = memory
        self.chat_history = [] 
        self.max_memory_turns = 10 # Keeps the last 10 interactions to avoid token overflow
        
        # 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
            self.log(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; self.log("WEB MODE")
            else:
                self.abs_path = os.path.abspath(data_source)
                ext = os.path.splitext(self.abs_path)[1].lower()
                
                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)
                    self.log(f"VISION MODE (Image: {os.path.basename(self.abs_path)})")
                    
                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)
                    self.log("SINGLE FILE MODE")
                    if ext in ['.csv', '.xlsx', '.xls', '.parquet', '.json']:
                        self.data_analyst_mode = True
                else: 
                    self.mode = "directory"; self.watch_target = self.abs_path; self.log("FOLDER MODE")
                    
        elif isinstance(data_source, (pd.DataFrame, pd.Series)) or (pl and isinstance(data_source, (pl.DataFrame, pl.Series))):
            self.mode = "memory"; self.memory_object = data_source; self.log("DATAFRAME MODE (Pandas/Polars)")
            self.data_analyst_mode = True
            self.dataframe_context = data_source
            
        elif isinstance(data_source, Callable):
            self.mode = "callable"; self.memory_object = data_source; self.log("CALLABLE DATAFRAME MODE (Dynamic)")
            self.data_analyst_mode = True
            self.dataframe_context = data_source()
            
        else:
            self.mode = "memory"; self.memory_object = data_source; self.log("MEMORY MODE")

        if reset_db: VectorDBFactory.reset("./db_storage", logger=self.log)

        self.embed_manager = EmbeddingsManager(model_name=embedding_model, source=embedding_source, logger=self.log)
        self.embedding_function = self.embed_manager.get_function()
        self.loader = UniversalLoader(logger=self.log)
        self.db = VectorDBFactory.create(vector_db, self.embedding_function, logger=self.log)
        
        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,
            logger=self.log,
            **llm_kwargs
        )
        
        self._initial_ingest()
        
        if auto_update 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, logger=self.log)
                self.watcher.start()
            elif self.mode == "callable":
                self.watcher = PollingWatcher(self._on_polling_event, self.memory_object, interval=poll_interval, logger=self.log)
                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, logger=self.log)
                t = threading.Thread(target=self.watcher.start); t.daemon = True; t.start()

    def _setup_directories_and_logging(self):
        """Creates public folders for graphs and logs."""
        os.makedirs("public/graph", exist_ok=True)
        os.makedirs("public/logs", exist_ok=True)
        
        self.logger = logging.getLogger("RostaingAgent")
        self.logger.setLevel(logging.INFO)
        self.logger.handlers = [] # Clear existing handlers
        
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        
        if self.save_logs:
            log_file = f"public/logs/agent_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
            fh = logging.FileHandler(log_file)
            fh.setFormatter(formatter)
            self.logger.addHandler(fh)
            
        ch = logging.StreamHandler()
        ch.setFormatter(formatter)
        self.logger.addHandler(ch)

    def log(self, message, level="info"):
        """Internal logging handler."""
        if level == "info": self.logger.info(message)
        elif level == "error": self.logger.error(message)
        elif level == "warning": self.logger.warning(message)

    def _initial_ingest(self):
        self.log("Starting data ingestion...")
        if self.mode == "database": self._process_database(self.db_config)
        elif self.mode == "web": self._process_web(self.web_url)
        elif self.mode in ["memory", "callable"]: self._ingest_memory_object()
        elif self.vision_mode: pass 
        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 c.get("type") == "mongodb"
                else (self.loader.process_neo4j(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):
        obj = self.memory_object() if self.mode == "callable" else self.memory_object
        docs = self.loader.process_dataframe_object(obj); self._update_db(docs)

    def _update_db(self, docs):
        if not docs: return
        # Big Data Optimization: dynamic chunk sizing based on doc count could be added here
        splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=250)
        chunks = splitter.split_documents(docs)
        self.db = VectorDBFactory.add_docs(self.db, chunks, self.db_type, self.embedding_function, "./db_storage", logger=self.log)
        self.log(f"RAG Knowledge Base updated ({len(chunks)} chunks).")

    def _on_polling_event(self, source): 
        if self.mode == "database": self._process_database(source)
        elif self.mode == "callable":
            self.dataframe_context = source()
            self._ingest_memory_object()
        else: self._process_web(source)
        
    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: 
            self.log(f"Processing file: {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 Exception as e: 
            self.log(f"Failed to process {file_path}: {e}", "error")

    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: 
                # MMR for diversity, fetching more to handle complex queries
                retriever = self.db.as_retriever(search_type="mmr", search_kwargs={"k": 8, "fetch_k": 30})
                docs = retriever.invoke(query)
        except Exception as e: 
            self.log(f"Retrieval Error: {e}", "error")
        return docs

    def chat(self, user_query):
        """Unified Agent Chat Endpoint with Autonomous Reasoning."""
        effective_stream = self.config['stream']
        output_fmt = self.config['output_format']
        
        # 1. RAG Retrieval
        context_text = ""
        if not self.vision_mode:
            docs = self._safe_retrieval(user_query)
            context_text = "\n\n".join([f"Source: {d.metadata.get('source', 'Unknown')}\nContent: {d.page_content}" for d in docs]) if docs else ""
            if context_text: self.log(f"RAG Context fetched: {len(docs)} segments used.")
        
        # 2. ChatGPT-like Sliding Window Memory
        history_str = ""
        if self.use_memory and self.chat_history:
            history_str = "CHAT HISTORY (Prioritize this for context flow):\n" + "\n".join([f"User: {q}\nAI: {a}" for q,a in self.chat_history])
        
        # 3. Execution Routing
        if self.data_analyst_mode and self.dataframe_context is not None:
            self.log("Engaging Autonomous Data Analyst Agent...")
            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:
            self.log("Engaging Reasoning & Knowledge Agent...")
            full_context = f"{history_str}\n\nKNOWLEDGE BASE CONTEXT (Strict Grounding):\n{context_text}"
            response = self.engine.generate(
                prompt=user_query,
                context=full_context,
                context_config=self.config,
                image_path=self.active_image_path,
                stream=effective_stream,
                output_format=output_fmt
            )

        # Update Memory intelligently
        if self.use_memory and not effective_stream: 
            self.chat_history.append((user_query, response))
            if len(self.chat_history) > self.max_memory_turns:
                self.chat_history.pop(0) # Sliding window
                
        return response