# Copyright (c) ModelScope Contributors. All rights reserved.
"""
Unified API endpoints for chat and search functionality
Provides WebSocket endpoint for real-time chat conversations with integrated search
"""
import platform
import time

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, HTTPException
from typing import Dict, Any, List, Optional, Union
from pydantic import BaseModel
import json
import asyncio
import uuid
from datetime import datetime
import random
import os
import threading
from sirchmunk.search import AgenticSearch
from sirchmunk.llm.openai_chat import OpenAIChat
from sirchmunk.api.components.history_storage import HistoryStorage
from sirchmunk.api.components.settings_storage import SettingsStorage
from sirchmunk.api.components.monitor_tracker import llm_usage_tracker


# Try to import tkinter for file dialogs
try:
    import tkinter as tk
    from tkinter import filedialog, messagebox
    TKINTER_AVAILABLE = True
except ImportError:
    TKINTER_AVAILABLE = False

router = APIRouter(prefix="/api/v1", tags=["chat", "search"])

# Initialize persistent history storage
history_storage = HistoryStorage()

# Initialize settings storage for LLM configuration (with error handling)
try:
    settings_storage = SettingsStorage()
except Exception as e:
    print(f"[WARNING] Failed to initialize SettingsStorage: {e}")
    settings_storage = None

# In-memory cache for active sessions (for backward compatibility)
chat_sessions = {}

# Active WebSocket connections
class ChatConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

# Unified log callback management
class WebSocketLogger:
    """
    WebSocket-aware logger that wraps websocket communications.
    
    Provides logger-style methods (info, warning, etc.) similar to loguru,
    with support for flush and end parameters for streaming output.
    Compatible with sirchmunk.utils.log_utils.AsyncLogger interface.
    """
    
    def __init__(self, websocket: WebSocket, manager: Optional[ChatConnectionManager] = None, log_type: str = "log", task_id: Optional[str] = None):
        """
        Initialize WebSocket logger.
        
        Args:
            websocket: WebSocket connection to send logs to
            manager: Optional ConnectionManager for routing messages
            log_type: Type of log message ("log" or "search_log")
            task_id: Optional task ID for grouping related log messages
        """
        self.websocket = websocket
        self.manager = manager
        self.log_type = log_type
        self.task_id = task_id or str(uuid.uuid4())  # Generate unique task ID
    
    async def _send_log(self, level: str, message: str, flush: bool = False, end: str = "\n"):
        """
        Send log message through WebSocket.
        
        Args:
            level: Log level (info, warning, error, etc.)
            message: Message content
            flush: If True, force immediate output (adds small delay for streaming)
            end: String appended after message (default: "\n")
        """
        # Append end character to message
        full_message = message + end if end else message
        
        # Determine if this is a streaming message (no timestamp prefix should be added on frontend)
        # Streaming condition: message should be appended to current line (end is not a newline)
        # This indicates it's part of a multi-chunk streaming output (like LLM responses)
        is_streaming = end != "\n"
        
        # Prepare log message
        log_data = {
            "type": self.log_type,
            "level": level,
            "message": full_message,
            "timestamp": datetime.now().isoformat(),
            "is_streaming": is_streaming,  # Flag for frontend to know if this is streaming output
            "task_id": self.task_id,  # Task ID for grouping related messages
            "flush": flush,  # Include flush flag for frontend handling
        }
        
        # Send through WebSocket
        if self.manager:
            await self.manager.send_personal_message(json.dumps(log_data), self.websocket)
        else:
            await self.websocket.send_text(json.dumps(log_data))
        
        # If flush is requested, add small delay for proper streaming
        if flush:
            await asyncio.sleep(0.01)  # Very short delay for streaming (reduced from 0.05s)
        else:
            await asyncio.sleep(0.05)  # Standard delay (reduced from 0.1s)
    
    async def log(self, level: str, message: str, flush: bool = False, end: str = "\n"):
        """Log a message at the specified level"""
        await self._send_log(level, message, flush=flush, end=end)
    
    async def debug(self, message: str, flush: bool = False, end: str = "\n"):
        """Log a debug message"""
        await self._send_log("debug", message, flush=flush, end=end)
    
    async def info(self, message: str, flush: bool = False, end: str = "\n"):
        """Log an info message"""
        await self._send_log("info", message, flush=flush, end=end)
    
    async def warning(self, message: str, flush: bool = False, end: str = "\n"):
        """Log a warning message"""
        await self._send_log("warning", message, flush=flush, end=end)
    
    async def error(self, message: str, flush: bool = False, end: str = "\n"):
        """Log an error message"""
        await self._send_log("error", message, flush=flush, end=end)
    
    async def success(self, message: str, flush: bool = False, end: str = "\n"):
        """Log a success message"""
        await self._send_log("success", message, flush=flush, end=end)
    
    async def critical(self, message: str, flush: bool = False, end: str = "\n"):
        """Log a critical message"""
        await self._send_log("critical", message, flush=flush, end=end)


class LogCallbackManager:
    """
    Centralized management for all log callback functions.
    
    Creates callback functions and logger instances that are compatible with
    sirchmunk.utils.log_utils.AsyncLogger interface, supporting flush and end parameters.
    """

    @staticmethod
    async def create_search_log_callback(websocket: WebSocket, manager: ChatConnectionManager, task_id: Optional[str] = None):
        """
        Create search log callback for chat WebSocket.
        
        Returns a callback function compatible with log_utils signature:
        async def callback(level: str, message: str, end: str, flush: bool)
        
        NOTE: The signature MUST match log_utils.LogCallback exactly:
              (level: str, message: str, end: str, flush: bool) -> None
        
        Args:
            websocket: WebSocket connection
            manager: Connection manager for routing
            task_id: Optional task ID for grouping related messages (auto-generated if not provided)
            
        Returns:
            Async callback function
        """
        # Generate unique task ID for this search session
        if task_id is None:
            task_id = f"search_{uuid.uuid4().hex[:8]}"
        
        logger = WebSocketLogger(websocket, manager, log_type="search_log", task_id=task_id)
        
        # Track recent messages for deduplication (message -> timestamp)
        recent_messages: Dict[str, float] = {}
        DEDUP_WINDOW_SEC = 0.5  # Messages within this window are considered duplicates

        # CRITICAL: This callback signature MUST match log_utils.LogCallback
        # Signature: (level: str, message: str, end: str, flush: bool) -> None
        async def search_log_callback(level: str, message: str, end: str, flush: bool):
            """
            Log callback compatible with log_utils.LogCallback type.
            
            Args:
                level: Log level (info, warning, error, etc.)
                message: Message content (WITHOUT end character appended)
                end: String to append after message
                flush: Whether to flush immediately
            """
            import time
            nonlocal recent_messages
            
            # Create unique key for this message (include level and message content)
            msg_key = f"{level}:{message}"
            current_time = time.time()
            
            # Check for duplicate within time window
            if msg_key in recent_messages:
                last_time = recent_messages[msg_key]
                if current_time - last_time < DEDUP_WINDOW_SEC:
                    # Skip duplicate message within dedup window
                    return
            
            # Clean up old entries (older than 2x window)
            cutoff = current_time - (DEDUP_WINDOW_SEC * 2)
            recent_messages = {k: v for k, v in recent_messages.items() if v > cutoff}
            
            # Record this message
            recent_messages[msg_key] = current_time
            
            await logger._send_log(level, message, flush=flush, end=end)
        
        return search_log_callback

    @staticmethod
    def create_logger(websocket: WebSocket, manager: Optional[ChatConnectionManager] = None, log_type: str = "log", task_id: Optional[str] = None) -> WebSocketLogger:
        """
        Create a WebSocketLogger instance with logger-style methods.
        
        This provides a logger interface similar to create_logger from log_utils,
        allowing usage like: await logger.info("message", flush=True, end="")
        
        Args:
            websocket: WebSocket connection
            manager: Optional ConnectionManager for routing messages
            log_type: Type of log message ("log" or "search_log")
            task_id: Optional task ID for grouping related messages (auto-generated if not provided)
            
        Returns:
            WebSocketLogger instance
            
        Example:
            logger = LogCallbackManager.create_logger(websocket, manager, "search_log")
            await logger.info("Processing started")
            await logger.info("Loading", flush=True, end=" -> ")
            await logger.success("Done!", flush=True)
        """
        if task_id is None:
            task_id = f"logger_{uuid.uuid4().hex[:8]}"
        return WebSocketLogger(websocket, manager, log_type, task_id)

manager = ChatConnectionManager()

# Search-related models and functions
class SearchRequest(BaseModel):
    query: str
    paths: Union[str, List[str]]  # Expects absolute file/directory paths from user's local filesystem
    mode: Optional[str] = "DEEP"
    max_depth: Optional[int] = 5
    top_k_files: Optional[int] = 3


def get_envs() -> Dict[str, Any]:
    """
    Get LLM configuration with priority resolution.

    Priority: SettingsStorage (WebUI) > os.getenv() (includes .env) > defaults
    """
    # Try to get from settings storage first (highest priority: user WebUI overrides)
    if settings_storage is not None:
        base_url = settings_storage.get_env_variable("LLM_BASE_URL", "")
        api_key = settings_storage.get_env_variable("LLM_API_KEY", "")
        model_name = settings_storage.get_env_variable("LLM_MODEL_NAME", "")
    else:
        base_url = ""
        api_key = ""
        model_name = ""

    # Fallback to os.getenv() which includes .env values loaded at startup
    if not base_url:
        base_url = os.getenv("LLM_BASE_URL", "https://api.openai.com/v1")
    if not api_key:
        api_key = os.getenv("LLM_API_KEY", "")
    if not model_name:
        model_name = os.getenv("LLM_MODEL_NAME", "gpt-5.2")

    print(f"[ENV CONFIG] base_url={base_url}, model_name={model_name}, api_key={'***' if api_key else '(not set)'}")

    return dict(
        base_url=base_url,
        api_key=api_key,
        model_name=model_name,
    )


def get_search_instance(log_callback=None):
    """
    Get configured search instance with optional log callback.
    
    Creates OpenAIChat instance with settings from DuckDB (priority) or environment variables (fallback).
    
    Args:
        log_callback: Optional callback for logging
        
    Returns:
        Configured AgenticSearch instance
    """
    # Get LLM configuration from settings storage (priority) or env variables (fallback)
    try:
        envs = get_envs()
        
        # Create OpenAI LLM instance with retrieved configuration
        llm = OpenAIChat(
            base_url=envs["base_url"],
            api_key=envs["api_key"],
            model=envs["model_name"],
            log_callback=log_callback,
        )
        
        # Create and return AgenticSearch instance with configured LLM
        return AgenticSearch(llm=llm, log_callback=log_callback)
    
    except Exception as e:
        # If settings retrieval fails, fall back to default AgenticSearch initialization
        print(f"[WARNING] Please config ENVs: LLM_BASE_URL, LLM_API_KEY, LLM_MODEL_NAME. Error: {e}")
        return AgenticSearch(log_callback=log_callback)


_COOLDOWN_SECONDS = 1.0
_DIALOG_LOCK = threading.Lock()
_LAST_CLOSE_TIME = 0
_ROOT_INSTANCE = None


def _get_bg_root():
    """
    Retrieves the global root window.
    Initializes it only once (Singleton pattern) to prevent lag.
    """
    global _ROOT_INSTANCE

    if threading.current_thread() is not threading.main_thread():
        raise RuntimeError("Tkinter must be executed on the Main Thread.")

    if _ROOT_INSTANCE is None or not _ROOT_INSTANCE.winfo_exists():
        _ROOT_INSTANCE = tk.Tk()
        _ROOT_INSTANCE.title("File Picker")
        _ROOT_INSTANCE.attributes("-alpha", 0.0)
        _ROOT_INSTANCE.withdraw()

    return _ROOT_INSTANCE


def open_file_dialog(dialog_type: str = "files", multiple: bool = True) -> List[str]:
    """
    Opens a native file picker dialog using tkinter.
    """
    global _LAST_CLOSE_TIME

    if not _DIALOG_LOCK.acquire(blocking=False):
        return []

    selected_paths = []

    try:
        if time.time() - _LAST_CLOSE_TIME < _COOLDOWN_SECONDS:
            return []

        root = _get_bg_root()
        root.deiconify()
        root.attributes("-topmost", True)
        root.lift()
        root.focus_force()

        if platform.system() == "Darwin":
            root.update_idletasks()
        else:
            root.update()

        kwargs = {"parent": root, "title": "Select File(s)"}

        # Set file types filter
        if dialog_type == "files":
            filetypes = [
                ("All Files", "*.*"),
                ("PDF Documents", "*.pdf"),
                ("Word Documents", "*.docx *.doc"),
                ("Excel Spreadsheets", "*.xlsx *.xls *.csv"),
                ("Images", "*.png *.jpg *.jpeg *.gif *.svg"),
                ("Text Files", "*.txt *.md *.json *.xml"),
            ]

            if multiple:
                res = filedialog.askopenfilenames(filetypes=filetypes, **kwargs)
                selected_paths = list(res) if res else []
            else:
                res = filedialog.askopenfilename(filetypes=filetypes, **kwargs)
                selected_paths = [res] if res else []

        elif dialog_type == "directory":
            kwargs["title"] = "Select Directory"
            res = filedialog.askdirectory(**kwargs)
            selected_paths = [res] if res else []

    except Exception as e:
        print(f"Dialog Error: {e}")
        selected_paths = []

    finally:
        if _ROOT_INSTANCE is not None and _ROOT_INSTANCE.winfo_exists():
            _ROOT_INSTANCE.attributes("-topmost", False)
            _ROOT_INSTANCE.withdraw()
            _ROOT_INSTANCE.update()

        _LAST_CLOSE_TIME = time.time()
        _DIALOG_LOCK.release()

    return selected_paths


async def _perform_web_search(query: str, websocket: WebSocket, manager: ChatConnectionManager) -> Dict[str, Any]:
    """
    Mock web search functionality
    TODO: Replace with actual web search implementation
    """
    await manager.send_personal_message(json.dumps({
        "type": "search_log",
        "level": "info",
        "message": "🌐 Starting web search...",
        "timestamp": datetime.now().isoformat()
    }), websocket)
    
    # Simulate web search delay
    await asyncio.sleep(random.uniform(0.5, 1.0))
    
    await manager.send_personal_message(json.dumps({
        "type": "search_log",
        "level": "info",
        "message": f"🔎 Searching web for: {query}",
        "timestamp": datetime.now().isoformat()
    }), websocket)
    
    await asyncio.sleep(random.uniform(0.5, 1.0))
    
    # Mock web search results
    web_results = {
        "sources": [
            {
                "url": "https://example.com/article1",
                "title": "Comprehensive Guide to " + query[:30],
                "snippet": "This article provides detailed information about the subject matter...",
                "relevance_score": 0.95
            },
            {
                "url": "https://example.com/article2", 
                "title": "Advanced Concepts and Applications",
                "snippet": "Exploring advanced techniques and real-world applications...",
                "relevance_score": 0.87
            },
            {
                "url": "https://example.com/article3",
                "title": "Latest Research and Findings",
                "snippet": "Recent discoveries and innovations in this field...",
                "relevance_score": 0.82
            }
        ],
        "summary": f"Found 3 relevant web sources for '{query}'. The sources cover comprehensive guides, advanced concepts, and latest research."
    }
    
    await manager.send_personal_message(json.dumps({
        "type": "search_log",
        "level": "success",
        "message": f"✅ Web search completed: found {len(web_results['sources'])} sources",
        "timestamp": datetime.now().isoformat()
    }), websocket)
    
    return web_results

async def _chat_only(
    message: str,
    websocket: WebSocket,
    manager: ChatConnectionManager
) -> tuple[str, Dict[str, Any]]:
    """
    Mode 1: Pure chat mode (no RAG, no web search)
    Direct LLM chat without any retrieval augmentation
    """
    try:
        await manager.send_personal_message(json.dumps({
            "type": "status",
            "stage": "generating",
            "message": "💬 Generating response..."
        }), websocket)
        
        # Create log callback for streaming LLM output
        llm_log_callback = await LogCallbackManager.create_search_log_callback(websocket, manager)
        
        # Initialize OpenAI client with log callback for streaming
        envs: Dict[str, Any] = get_envs()
        llm = OpenAIChat(
            api_key=envs["api_key"],
            base_url=envs["base_url"],
            model=envs["model_name"],
            log_callback=llm_log_callback
        )
        
        # Prepare messages for LLM
        messages = [
            {"role": "system", "content": "You are a helpful AI assistant. Provide clear, accurate, and helpful responses."},
            {"role": "user", "content": message}
        ]
        
        # Generate response with streaming
        llm_response = await llm.achat(messages=messages, stream=True)
        
        # Record LLM usage for monitoring (always record call, even if usage is empty)
        # Some LLM APIs don't return usage in streaming mode
        usage_data = llm_response.usage if llm_response.usage else {}
        llm_usage_tracker.record_usage(
            model=llm_response.model or envs["model_name"],
            usage=usage_data
        )
        
        sources = {}
        
        return llm_response.content, sources
    
    except Exception as e:
        # Send error message to frontend
        await manager.send_personal_message(json.dumps({
            "type": "error",
            "message": f"LLM chat failed: {str(e)}"
        }), websocket)
        
        # Re-raise to be caught by outer handler
        raise


async def _chat_rag(
    message: str,
    kb_name: str,
    websocket: WebSocket,
    manager: ChatConnectionManager
) -> tuple[str, Dict[str, Any]]:
    """
    Mode 2: Chat + RAG (enable_rag=True, enable_web_search=False)
    LLM chat with knowledge base retrieval
    """
    sources = {}
    if not kb_name:
        await manager.send_personal_message(json.dumps({
            "type": "error",
            "message": "No search paths specified for RAG search."
        }), websocket)
        response = "Please specify search paths for RAG search."
        return response, sources
    
    try:
        # Create log callback for streaming search logs
        search_log_callback = await LogCallbackManager.create_search_log_callback(websocket, manager)

        # Create search instance with log callback
        search_engine = get_search_instance(log_callback=search_log_callback)

        paths = [path.strip() for path in kb_name.split(",")]
        await search_log_callback("info", f"📂 Parsed search paths: {paths}", "\n", False)

        # Execute RAG search
        print(f"[MODE 2] RAG search with query: {message}, paths: {paths}")
        
        search_result = await search_engine.search(
            query=message,
            paths=paths,
            top_k_files=3,
        )

        # Calculate the llm usage
        for usage in search_engine.llm_usages:
            llm_usage_tracker.record_usage(
                model=search_engine.llm._model,
                usage=usage,
            )

        # Send search completion
        await manager.send_personal_message(json.dumps({
            "type": "search_complete",
            "message": "✅ Knowledge base search completed"
        }), websocket)
        
        # Use search result as response
        response = search_result
        
        # Add RAG sources
        sources["rag"] = [
            {
                "kb_name": kb_name,
                "content": f"Retrieved content from {kb_name}",
                "relevance_score": 0.92
            }
        ]
        
    except Exception as e:
        # Send search error
        await manager.send_personal_message(json.dumps({
            "type": "search_error",
            "message": f"❌ RAG search failed: {str(e)}"
        }), websocket)
        
        # Fallback to chat only
        await manager.send_personal_message(json.dumps({
            "type": "status",
            "stage": "fallback",
            "message": "⚠️ RAG mode did not find relevant results, falling back to pure chat mode..."
        }), websocket)
        
        print(f"[MODE 2] RAG search failed, falling back to chat only: {str(e)}")
        response, sources = await _chat_only(message, websocket, manager)
    
    return response, sources


async def _chat_web_search(
    message: str,
    websocket: WebSocket,
    manager: ChatConnectionManager
) -> tuple[str, Dict[str, Any]]:
    """
    Mode 3: Chat + Web Search (enable_rag=False, enable_web_search=True)
    LLM chat with web search augmentation (currently mock)
    """
    await manager.send_personal_message(json.dumps({
        "type": "status",
        "stage": "web_search",
        "message": "🌐 Searching the web..."
    }), websocket)
    
    # Perform mock web search
    web_results = await _perform_web_search(message, websocket, manager)
    
    # Check if web search returned valid results
    if not web_results or not web_results.get("sources"):
        # Fallback to chat only
        await manager.send_personal_message(json.dumps({
            "type": "status",
            "stage": "fallback",
            "message": "⚠️ Web search did not return results, falling back to pure chat mode..."
        }), websocket)
        
        print(f"[MODE 3] Web search failed, falling back to chat only")
        response, sources = await _chat_only(message, websocket, manager)
        return response, sources
    
    # Generate response enhanced with web search results
    web_context = "\n\nBased on web search results:\n"
    for source in web_results["sources"]:
        web_context += f"- {source['title']}: {source['snippet']}\n"
    
    # Use LLM to generate response with web context
    await manager.send_personal_message(json.dumps({
        "type": "status",
        "stage": "generating",
        "message": "💬 Generating response with web context..."
    }), websocket)

    envs: Dict[str, Any] = get_envs()
    llm_log_callback = await LogCallbackManager.create_search_log_callback(websocket, manager)
    llm = OpenAIChat(
        api_key=envs["api_key"],
        base_url=envs["base_url"],
        model=envs["model_name"],
        log_callback=llm_log_callback
    )
    
    messages = [
        {"role": "system", "content": "You are a helpful AI assistant. Use the provided web search results to answer the user's question accurately."},
        {"role": "user", "content": f"{message}\n\nWeb search context:\n{web_context}"}
    ]
    
    llm_response = await llm.achat(messages=messages, stream=True)
    
    # Record LLM usage for monitoring (always record call, even if usage is empty)
    usage_data = llm_response.usage if llm_response.usage else {}
    llm_usage_tracker.record_usage(
        model=llm_response.model or envs["model_name"],
        usage=usage_data
    )
    
    sources = {"web": web_results["sources"]}
    
    return llm_response.content, sources


async def _chat_rag_web_search(
    message: str,
    kb_name: str,
    websocket: WebSocket,
    manager: ChatConnectionManager
) -> tuple[str, Dict[str, Any]]:
    """
    Mode 4: Chat + RAG + Web Search (enable_rag=True, enable_web_search=True)
    LLM chat with both knowledge base retrieval and web search
    """
    sources = {}
    if not kb_name:
        await manager.send_personal_message(json.dumps({
            "type": "error",
            "message": "No search paths specified for RAG search."
        }), websocket)
        response = "Please specify search paths for RAG search."
        return response, sources

    # Step 1: Perform RAG search
    try:
        search_log_callback = await LogCallbackManager.create_search_log_callback(websocket, manager)

        search_engine = get_search_instance(log_callback=search_log_callback)
        paths = [path.strip() for path in kb_name.split(",")]
        await search_log_callback("info", f"📂 RAG search paths: {paths}", "\n", False)

        print(f"[MODE 4] RAG search with query: {message}, paths: {paths}")
        
        rag_result = await search_engine.search(
            query=message,
            paths=paths,
            top_k_files=3,
        )

        for rag_usage in search_engine.llm_usages:
            llm_usage_tracker.record_usage(
                model=search_engine.llm._model,
                usage=rag_usage,
            )
        
        await manager.send_personal_message(json.dumps({
            "type": "search_complete",
            "message": "✅ Knowledge base search completed"
        }), websocket)
        
        sources["rag"] = [
            {
                "kb_name": kb_name,
                "content": f"Retrieved from {kb_name}",
                "relevance_score": 0.92
            }
        ]
        
    except Exception as e:
        await manager.send_personal_message(json.dumps({
            "type": "search_error",
            "message": f"⚠️ RAG search failed: {str(e)}, continuing with web search..."
        }), websocket)
        rag_result = f"[RAG search unavailable: {str(e)}]"
        sources["rag"] = [{"error": str(e)}]
    
    # Step 2: Perform web search
    await manager.send_personal_message(json.dumps({
        "type": "status",
        "stage": "web_search",
        "message": "🌐 Step 2/2: Searching the web..."
    }), websocket)

    # TODO: add llm usage
    web_results = await _perform_web_search(message, websocket, manager)
    sources["web"] = web_results["sources"]
    
    # Combine results
    web_context = "\n\n## Additional Web Sources:\n"
    for source in web_results["sources"]:
        web_context += f"- [{source['title']}]({source['url']})\n"
    
    # If RAG succeeded, use it as primary response; otherwise use web search only
    if rag_result and "[RAG search unavailable" not in rag_result:
        response = rag_result + web_context
    else:
        response = f"Based on web search results:\n{web_context}"
    
    return response, sources


# WebSocket endpoint for chat with integrated search
@router.websocket("/chat")
async def chat_websocket(websocket: WebSocket):
    """
    WebSocket endpoint for real-time chat conversations with integrated search
    
    Supports 4 modes:
    1. Pure chat: enable_rag=False, enable_web_search=False
    2. Chat + RAG: enable_rag=True, enable_web_search=False
    3. Chat + Web Search: enable_rag=False, enable_web_search=True (mock)
    4. Chat + RAG + Web Search: enable_rag=True, enable_web_search=True (RAG real, web mock)
    """
    await manager.connect(websocket)
    
    try:
        while True:
            # Receive message from client
            data = await websocket.receive_text()
            request_data = json.loads(data)
            
            message = request_data.get("message", "")
            session_id = request_data.get("session_id")
            history = request_data.get("history", [])
            kb_name = request_data.get("kb_name", "")
            enable_rag = request_data.get("enable_rag", False)
            enable_web_search = request_data.get("enable_web_search", False)

            print(f"\n{'='*60}")
            print(f"[CHAT REQUEST] Message: {message[:50]}...")
            print(f"[CHAT REQUEST] KB: {kb_name}, RAG: {enable_rag}, Web: {enable_web_search}")
            print(f"{'='*60}\n")
            
            # Generate or use existing session ID
            if not session_id:
                session_id = f"chat_{uuid.uuid4().hex[:8]}"
            
            # Send session ID to client
            await manager.send_personal_message(json.dumps({
                "type": "session",
                "session_id": session_id
            }), websocket)
            
            # Store session data (in-memory + persistent)
            if session_id not in chat_sessions:
                chat_sessions[session_id] = {
                    "session_id": session_id,
                    "title": f"Chat Session",
                    "messages": [],
                    "created_at": datetime.now().isoformat(),
                    "updated_at": datetime.now().isoformat(),
                    "settings": {
                        "kb_name": kb_name,
                        "enable_rag": enable_rag,
                        "enable_web_search": enable_web_search
                    }
                }
                # Save new session to persistent storage
                history_storage.save_session(chat_sessions[session_id])
            
            # Update session with new message
            session = chat_sessions[session_id]
            user_message = {
                "role": "user",
                "content": message,
                "timestamp": datetime.now().isoformat()
            }
            session["messages"].append(user_message)
            session["updated_at"] = datetime.now().isoformat()
            
            # Save user message to persistent storage
            history_storage.save_message(session_id, user_message)
            
            # ============================================================
            # Route to appropriate chat mode based on feature flags
            # ============================================================
            response = ""
            sources = {}
            
            if enable_rag and enable_web_search:
                # Mode 4: Chat + RAG + Web Search
                print(f"[MODE 4] Chat + RAG + Web Search")
                response, sources = await _chat_rag_web_search(
                    message, kb_name, websocket, manager
                )
                
            elif enable_rag and not enable_web_search:
                # Mode 2: Chat + RAG
                print(f"[MODE 2] Chat + RAG")
                response, sources = await _chat_rag(
                    message, kb_name, websocket, manager
                )
                    
            elif not enable_rag and enable_web_search:
                # Mode 3: Chat + Web Search only
                print(f"[MODE 3] Chat + Web Search only")
                response, sources = await _chat_web_search(
                    message, websocket, manager
                )
                
            else:
                # Mode 1: Pure chat (no RAG, no web search)
                print(f"[MODE 1] Pure chat mode")
                response, sources = await _chat_only(
                    message, websocket, manager
                )
            
            # ============================================================
            # Stream response to client
            # ============================================================
            words = response.split()
            
            for i, word in enumerate(words):
                await manager.send_personal_message(json.dumps({
                    "type": "stream",
                    "content": word + " "
                }), websocket)
                
                # Add small delay for realistic streaming
                if i % 3 == 0:  # Every 3 words
                    await asyncio.sleep(0.05)
            
            # Send sources if available
            if sources:
                await manager.send_personal_message(json.dumps({
                    "type": "sources",
                    **sources
                }), websocket)
            
            # Send final result
            await manager.send_personal_message(json.dumps({
                "type": "result",
                "content": response.strip(),
                "session_id": session_id
            }), websocket)
            
            # Store assistant response in session
            assistant_message = {
                "role": "assistant",
                "content": response.strip(),
                "sources": sources if sources else None,
                "timestamp": datetime.now().isoformat()
            }
            session["messages"].append(assistant_message)
            
            # Save assistant message to persistent storage
            history_storage.save_message(session_id, assistant_message)
            
            # Update session in persistent storage
            history_storage.save_session(session)
            
    except WebSocketDisconnect:
        manager.disconnect(websocket)
    except Exception as e:
        print(f"[ERROR] WebSocket error: {str(e)}")
        import traceback
        traceback.print_exc()
        try:
            await manager.send_personal_message(json.dumps({
                "type": "error",
                "message": f"An error occurred: {str(e)}"
            }), websocket)
        except:
            pass
        manager.disconnect(websocket)


# File picker endpoints
@router.post("/file-picker")
async def open_file_picker(request: Dict[str, Any]):
    """
    Open native file picker dialog using tkinter
    Returns real absolute paths from user's local filesystem
    """
    if not TKINTER_AVAILABLE:
        return {
            "success": False,
            "error": "Tkinter not available on this system",
            "data": []
        }
    
    dialog_type = request.get("type", "files")  # "files" or "directory"
    multiple = request.get("multiple", True)
    
    try:
        # Get absolute paths from user's local filesystem
        selected_paths = open_file_dialog(dialog_type, multiple)
        
        # Convert to absolute paths and validate they exist
        validated_paths = []
        for path in selected_paths:
            abs_path = os.path.abspath(path)
            if os.path.exists(abs_path):
                validated_paths.append(abs_path)
        
        return {
            "success": True,
            "data": {
                "paths": validated_paths,
                "count": len(validated_paths),
                "type": dialog_type,
                "multiple": multiple
            }
        }
        
    except Exception as e:
        return {
            "success": False,
            "error": f"Failed to open file picker: {str(e)}",
            "data": []
        }

@router.get("/file-picker/status")
async def get_file_picker_status():
    """Check if file picker is available on this system"""
    return {
        "success": True,
        "data": {
            "tkinter_available": TKINTER_AVAILABLE,
            "supported_types": ["files", "directory"] if TKINTER_AVAILABLE else [],
            "features": {
                "multiple_files": TKINTER_AVAILABLE,
                "directory_selection": TKINTER_AVAILABLE,
                "absolute_paths": TKINTER_AVAILABLE
            }
        }
    }

# Chat session management endpoints
@router.get("/chat/sessions")
async def get_chat_sessions(limit: int = 20, offset: int = 0):
    """Get list of chat sessions"""
    sessions_list = list(chat_sessions.values())
    # Sort by updated_at (most recent first)
    sessions_list.sort(key=lambda x: x["updated_at"], reverse=True)
    
    # Apply pagination
    paginated_sessions = sessions_list[offset:offset + limit]
    
    # Format for response
    formatted_sessions = []
    for session in paginated_sessions:
        last_message = ""
        if session["messages"]:
            last_msg = session["messages"][-1]
            last_message = last_msg["content"][:100] + "..." if len(last_msg["content"]) > 100 else last_msg["content"]
        
        formatted_sessions.append({
            "session_id": session["session_id"],
            "title": session.get("title", "Chat Session"),
            "message_count": len(session["messages"]),
            "last_message": last_message,
            "created_at": int(datetime.fromisoformat(session["created_at"]).timestamp()),
            "updated_at": int(datetime.fromisoformat(session["updated_at"]).timestamp()),
            "topics": ["AI", "Learning"]  # Mock topics
        })
    
    return {
        "success": True,
        "data": formatted_sessions,
        "pagination": {
            "limit": limit,
            "offset": offset,
            "total": len(sessions_list)
        }
    }

@router.get("/chat/sessions/{session_id}")
async def get_chat_session(session_id: str):
    """Get specific chat session details"""
    if session_id not in chat_sessions:
        raise HTTPException(status_code=404, detail="Chat session not found")
    
    session = chat_sessions[session_id]
    
    return {
        "success": True,
        "data": {
            "session_id": session["session_id"],
            "title": session.get("title", "Chat Session"),
            "messages": session["messages"],
            "settings": session.get("settings", {}),
            "created_at": session["created_at"],
            "updated_at": session["updated_at"]
        }
    }

@router.post("/chat/sessions/{session_id}/load")
async def load_chat_session(session_id: str):
    """Load chat session for continuation"""
    if session_id not in chat_sessions:
        raise HTTPException(status_code=404, detail="Chat session not found")
    
    session = chat_sessions[session_id]
    
    return {
        "success": True,
        "message": f"Chat session loaded successfully",
        "data": {
            "session_id": session_id,
            "title": session.get("title", "Chat Session"),
            "message_count": len(session["messages"]),
            "loaded_at": datetime.now().isoformat()
        }
    }

# Legacy search endpoints for backward compatibility
@router.get("/search/{kb_name}/suggestions")
async def get_search_suggestions(kb_name: str, query: str, limit: int = 8):
    """Get search suggestions - kept for backward compatibility"""
    # For now, return empty suggestions since we're using real file search
    if not query or len(query.strip()) < 2:
        return {
            "success": True,
            "data": [],
            "query": query
        }

    return {
        "success": True,
        "data": [],
        "query": query,
        "total_matches": 0
    }

@router.get("/search/knowledge-bases")
async def get_knowledge_bases():
    """Get list of available knowledge bases for search"""
    # Return empty list since we're using direct file paths now
    return {
        "success": True,
        "data": []
    }