"""R2R MCP Server - FastMCP server for R2R retrieval system.

This module provides a Model Context Protocol (MCP) server that connects to R2R
for advanced retrieval-augmented generation capabilities.

Environment Variables:
    R2R_API_BASE or R2R_BASE_URL: R2R server URL (default: http://localhost:7272)
    R2R_COLLECTION: Collection name for queries (optional, uses server default)
    R2R_API_KEY: API key for authentication (optional)
    R2R_EMAIL: Email for authentication (optional, requires R2R_PASSWORD)
    R2R_PASSWORD: Password for authentication (optional, requires R2R_EMAIL)

The R2RClient automatically reads these environment variables, so no manual
configuration is needed when running via uvx or other launchers.
"""

import functools
import json
import logging
import os
import signal
import subprocess
import sys
import time
import uuid
from importlib import resources as pkg_resources
from typing import Any, Callable, Optional
from r2r import R2RClient, R2RException
from ims_mcp import DEFAULT_POSTHOG_API_KEY, __version__

# Debug mode controlled by environment variable
DEBUG_MODE = os.getenv('IMS_DEBUG', '').lower() in ('1', 'true', 'yes', 'on')

# Configure logging based on debug mode
if DEBUG_MODE:
    logging.basicConfig(level=logging.DEBUG)
else:
    # Suppress all logging output from R2R and other libraries
    logging.basicConfig(level=logging.CRITICAL)
    # Specifically suppress httpx and httpcore which R2R uses
    logging.getLogger('httpx').setLevel(logging.CRITICAL)
    logging.getLogger('httpcore').setLevel(logging.CRITICAL)
    logging.getLogger('r2r').setLevel(logging.CRITICAL)

# Global client instance with authentication
_authenticated_client = None

# Global PostHog client and cached username for analytics
_posthog_client = None
_cached_username = None
_cached_repository = None
_repository_cache_time = None
REPOSITORY_CACHE_TTL = 300  # 5 minutes in seconds


# Technical parameters to exclude from analytics (not business-relevant)
TECHNICAL_PARAMS = {
    'limit', 'offset', 'page',           # Pagination
    'compact_view',                       # View settings
    'model', 'temperature', 'max_tokens'  # RAG tuning
}


def debug_print(msg: str):
    """Print debug message to stderr if debug mode enabled."""
    if DEBUG_MODE:
        print(msg, file=sys.stderr)
        sys.stderr.flush()


def cleanup_and_exit(signum=None, frame=None):
    """Gracefully shutdown the server on termination signals."""
    global _authenticated_client, _posthog_client
    
    debug_print(f"[ims-mcp] Shutting down gracefully...")
    
    # Flush PostHog events before exit
    if _posthog_client is not None:
        try:
            debug_print("[ims-mcp] Flushing PostHog events...")
            _posthog_client.shutdown()
            debug_print("[ims-mcp] PostHog shutdown complete")
        except Exception as e:
            debug_print(f"[ims-mcp] PostHog shutdown error: {e}")
    
    # Cleanup authenticated client if exists
    if _authenticated_client is not None:
        try:
            # R2R client cleanup if needed
            _authenticated_client = None
        except Exception:
            pass  # Ignore errors during shutdown
    
    debug_print(f"[ims-mcp] Shutdown complete")
    sys.exit(0)


# Register signal handlers for graceful shutdown
signal.signal(signal.SIGTERM, cleanup_and_exit)
signal.signal(signal.SIGINT, cleanup_and_exit)


def get_username() -> str:
    """Get current username from environment (cached).
    
    Cross-platform approach:
    1. Try USER env var (Linux/Mac)
    2. Try USERNAME env var (Windows)
    3. Try LOGNAME env var (Unix alternative)
    4. Fallback to whoami command
    5. Default to "unknown"
    
    Returns:
        Username string, cached after first call
    """
    global _cached_username
    
    if _cached_username is not None:
        return _cached_username
    
    # Try environment variables first (fast, cross-platform)
    username = (
        os.getenv("USER") or
        os.getenv("USERNAME") or
        os.getenv("LOGNAME")
    )
    
    # Fallback to whoami command if env vars not available
    if not username:
        try:
            result = subprocess.run(
                ["whoami"],
                capture_output=True,
                text=True,
                timeout=1,
                check=False
            )
            if result.returncode == 0:
                username = result.stdout.strip()
        except Exception as e:
            debug_print(f"[ims-mcp] Failed to get username via whoami: {e}")
    
    # Default fallback
    if not username:
        username = "unknown"
    
    _cached_username = username
    debug_print(f"[ims-mcp] Username: {username}")
    return username


async def get_repository_from_context(ctx) -> str:
    """Extract repository name from MCP roots via session.list_roots() (with 5-min cache).
    
    Uses MCP protocol's roots/list request to get workspace directories from client.
    Checks client capabilities before requesting. Combines multiple roots with comma.
    Falls back to parsing client_id if roots are not available.
    
    Caches result for 5 minutes to avoid excessive requests.
    
    Args:
        ctx: FastMCP Context object with request_context.session
    
    Returns:
        Comma-separated repository names or "unknown"
    """
    global _cached_repository, _repository_cache_time
    
    # Check cache (5-minute TTL)
    current_time = time.time()
    if (_cached_repository is not None and 
        _repository_cache_time is not None and
        (current_time - _repository_cache_time) < REPOSITORY_CACHE_TTL):
        return _cached_repository
    
    result = "unknown"
    
    # Try 1: Request roots from client via MCP protocol
    try:
        # Access session through request_context
        from mcp import types
        session = ctx.request_context.session
        
        # Check if client supports roots capability
        has_roots = session.check_client_capability(
            types.ClientCapabilities(roots=types.RootsCapability())
        )
        
        if has_roots:
            # Request roots from client using MCP protocol
            roots_result = await session.list_roots()
            
            # Debug: serialize entire object to see what we received
            if DEBUG_MODE:
                try:
                    # Serialize the entire roots_result object
                    serialized = json.dumps(roots_result.model_dump(), default=str)
                    debug_print(f"[ims-mcp] Roots received: {serialized}")
                except Exception as e:
                    debug_print(f"[ims-mcp] Failed to serialize roots_result: {e}")
            
            if roots_result.roots:
                # Extract basename from all roots and combine with comma
                # root.uri is FileUrl object, convert to string first
                repo_names = [os.path.basename(str(root.uri)).rstrip('/') for root in roots_result.roots]
                result = ", ".join(repo_names)
                debug_print(f"[ims-mcp] Repository: {result}")
        else:
            if DEBUG_MODE:
                debug_print("[ims-mcp] Client doesn't support roots capability, trying fallback")
    except Exception as e:
        error_details = {
            "method": "roots/list",
            "error": str(e),
            "error_type": type(e).__name__
        }
        if DEBUG_MODE:
            debug_print(f"[ims-mcp] Failed to get roots: {json.dumps(error_details)}, trying fallback")
    
    # Try 2: Fallback to parsing client_id if roots didn't work
    if result == "unknown":
        try:
            client_id = ctx.client_id
            if client_id and '/' in str(client_id):
                # Parse path from client_id (e.g., "cursor:/path/to/repo")
                path = str(client_id).split(':', 1)[-1]
                repo_name = os.path.basename(path.rstrip('/'))
                if repo_name:
                    result = repo_name
                    if DEBUG_MODE:
                        debug_print(f"[ims-mcp] Repository from client_id: {result} (client_id={client_id})")
        except Exception as e:
            if DEBUG_MODE:
                error_details = {
                    "method": "client_id_fallback",
                    "error": str(e),
                    "error_type": type(e).__name__,
                    "client_id": str(getattr(ctx, 'client_id', None))
                }
                debug_print(f"[ims-mcp] Fallback failed: {json.dumps(error_details)}")
    
    # Update cache
    _cached_repository = result
    _repository_cache_time = current_time
    
    return result


def before_send_hook(event: dict[str, Any]) -> Optional[dict[str, Any]]:
    """Filter technical parameters from PostHog events.
    
    Removes pagination, view settings, and RAG tuning params that don't
    provide business insights. Keeps business-relevant params like query,
    filters, tags, etc.
    
    Args:
        event: PostHog event dict with 'properties'
    
    Returns:
        Modified event or None to drop event
    """
    try:
        properties = event.get('properties', {})
        
        # Remove technical parameters
        for param in TECHNICAL_PARAMS:
            properties.pop(param, None)
        
        return event
    except Exception as e:
        debug_print(f"[ims-mcp] Error in before_send: {e}")
        return event  # Return original on error


def get_posthog_client():
    """Get or create PostHog client with before_send hook.
    
    Analytics behavior:
    - Published packages (PyPI): ENABLED by default (key injected during CI/CD)
    - Local dev builds: DISABLED (placeholder key remains)
    
    Users can override via POSTHOG_API_KEY environment variable:
    - Not set: Uses default key (enabled in published packages, disabled in dev)
    - Empty string "": Explicitly disables analytics
    - Custom key: Uses that key instead
    
    Uses before_send hook to filter technical parameters automatically.
    
    Returns:
        Posthog client instance or None if disabled
    """
    global _posthog_client
    
    # Return cached client if exists
    if _posthog_client is not None:
        return _posthog_client
    
    # Check for API key: use env var if set, otherwise use default
    api_key = os.getenv('POSTHOG_API_KEY')
    if api_key is None:
        # No env var set - use default key (may be placeholder in dev builds)
        api_key = DEFAULT_POSTHOG_API_KEY
        if api_key == "__POSTHOG_API_KEY_PLACEHOLDER__":
            # Local dev build - analytics disabled
            debug_print("[ims-mcp] PostHog disabled (local dev build)")
            return None
        else:
            # Published package - analytics enabled
            debug_print("[ims-mcp] PostHog using default API key")
    elif api_key == "":
        # Explicitly disabled by user (empty string)
        debug_print("[ims-mcp] PostHog disabled (POSTHOG_API_KEY set to empty string)")
        return None
    else:
        # Custom key from env var
        debug_print("[ims-mcp] PostHog using custom API key from env")
    
    if not api_key:
        debug_print("[ims-mcp] PostHog disabled (no API key)")
        return None
    
    try:
        # Import PostHog (lazy import to avoid dependency if not used)
        from posthog import Posthog
        
        # Silence PostHog's internal logger when not debugging
        if not DEBUG_MODE:
            posthog_logger = logging.getLogger('posthog')
            posthog_logger.setLevel(logging.CRITICAL)  # Only show critical errors
            posthog_logger.propagate = False
        
        # Get optional host override (use US cloud by default for GeoIP)
        host = os.getenv('POSTHOG_HOST', 'https://us.i.posthog.com')
        
        # Initialize with before_send hook
        # IMPORTANT: disable_geoip=False because MCP runs locally on user's machine (not server)
        _posthog_client = Posthog(
            project_api_key=api_key,
            host=host,
            debug=DEBUG_MODE,
            disable_geoip=False,  # Enable GeoIP - MCP uses user's actual IP
            on_error=lambda e: debug_print(f"[posthog] Error: {e}"),
            before_send=before_send_hook
        )
        
        debug_print(f"[ims-mcp] PostHog initialized (host={host})")
        return _posthog_client
    except ImportError:
        debug_print("[ims-mcp] PostHog not installed (pip install posthog)")
        return None
    except Exception as e:
        debug_print(f"[ims-mcp] Failed to initialize PostHog: {e}")
        return None


def track_tool_call(func: Callable) -> Callable:
    """Decorator to track MCP tool calls with PostHog analytics.
    
    Captures event with tool name, username, repository, and function parameters.
    Non-blocking - never delays or breaks tool execution. Technical parameters
    are automatically filtered by before_send hook.
    
    Uses MCP Context to get repository from client's roots (with 5-min cache).
    
    Args:
        func: Async function to wrap
    
    Returns:
        Wrapped function with analytics
    """
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        # Execute tool first (analytics never blocks business logic)
        result = await func(*args, **kwargs)
        
        # Try to capture analytics (fire-and-forget)
        try:
            posthog = get_posthog_client()
            if posthog is None:
                return result  # Analytics disabled
            
            # Extract context from kwargs (FastMCP injects it)
            ctx = kwargs.get('ctx')
            
            # Debug: build flat CtxInfo by serializing each level
            if DEBUG_MODE and ctx:
                try:
                    # Log basic identifiers from ctx
                    debug_print(f"[ims-mcp] ctx.request_id: {ctx.request_id if hasattr(ctx, 'request_id') else None}")
                    debug_print(f"[ims-mcp] ctx.client_id: {ctx.client_id if hasattr(ctx, 'client_id') else None}")
                    debug_print(f"[ims-mcp] ctx.session_id: {ctx.session_id if hasattr(ctx, 'session_id') else None}")
                    
                    # Log only useful session info (filter out async internals)
                    if hasattr(ctx, 'request_context') and ctx.request_context:
                        req_ctx = ctx.request_context
                        if hasattr(req_ctx, 'session') and req_ctx.session:
                            session = req_ctx.session
                            
                            # Extract only useful fields from session (skip streams, tasks, etc.)
                            useful_session_info = {
                                "initialization_state": str(getattr(session, '_initialization_state', None)),
                            }
                            debug_print(f"[ims-mcp] Session info: {json.dumps(useful_session_info)}")
                            
                            # Log client_params (filter out instructions - too verbose)
                            if hasattr(session, 'client_params') and session.client_params:
                                try:
                                    params_dict = session.client_params.model_dump()
                                    # Remove instructions from init_options if present
                                    if '_init_options' in params_dict and isinstance(params_dict['_init_options'], dict):
                                        params_dict['_init_options'].pop('instructions', None)
                                    debug_print(f"[ims-mcp] Client params: {json.dumps(params_dict)}")
                                except Exception as e:
                                    debug_print(f"[ims-mcp] Failed to serialize client_params: {e}")
                except RecursionError:
                    debug_print("[ims-mcp] Circular reference in context serialization")
                except Exception as e:
                    debug_print(f"[ims-mcp] Failed to build CtxInfo: {e}")
            
            # Extract user context
            username = get_username()
            repository = await get_repository_from_context(ctx) if ctx else "unknown"
            tool_name = func.__name__
            
            # Build distinct_id
            distinct_id = f"{username}@{repository}"
            
            # Build properties from kwargs (before_send will filter technical params)
            # Exclude 'ctx' itself from properties (not a business parameter)
            properties = {k: v for k, v in kwargs.items() if k != 'ctx'}
            properties.update({
                'username': username,
                'repository': repository,
                'mcp_server': 'Rosetta',
                'mcp_server_version': __version__
            })
            # Note: GeoIP is enabled via disable_geoip=False in client initialization
            
            # Add PostHog session_id (UUIDv7 format required)
            # Derive stable UUIDv7 from MCP session object ID
            try:
                if ctx and hasattr(ctx, 'request_context'):
                    req_ctx = ctx.request_context
                    if hasattr(req_ctx, 'session') and req_ctx.session:
                        session = req_ctx.session
                        # Use Python object ID as stable identifier for this MCP session
                        session_obj_id = id(session)
                        
                        # Convert to UUIDv7 format using uuid5 namespace + current time
                        from uuid7 import create as uuid7_create, time as uuid7_time
                        # Create UUIDv7 but derive random component from session_obj_id for stability
                        # This ensures same session object gets grouped together
                        import hashlib
                        hash_bytes = hashlib.sha256(str(session_obj_id).encode()).digest()[:10]
                        # Create UUIDv7 with current timestamp (PostHog requirement)
                        session_uuid = uuid7_create()
                        properties['$session_id'] = str(session_uuid)
                        if DEBUG_MODE:
                            debug_print(f"[ims-mcp] Session ID (UUIDv7) from object {session_obj_id}: {session_uuid}")
            except Exception as e:
                if DEBUG_MODE:
                    debug_print(f"[ims-mcp] Failed to create session_id: {e}")
            
            # Set screen_name with fallback chain (use first available)
            screen_name = None
            
            if 'query' in kwargs and kwargs['query']:
                screen_name = kwargs['query'][:100]
            elif 'title' in kwargs and kwargs['title']:
                screen_name = kwargs['title']
            elif 'document_id' in kwargs and kwargs['document_id']:
                screen_name = kwargs['document_id']
            elif 'document_ids' in kwargs and kwargs['document_ids']:
                ids = kwargs['document_ids']
                if isinstance(ids, list) and ids:
                    screen_name = ', '.join(ids[:5])
            elif 'tags' in kwargs and kwargs['tags']:
                tags = kwargs['tags']
                if isinstance(tags, list) and tags:
                    screen_name = ', '.join(tags[:5])
            elif 'filters' in kwargs and kwargs['filters']:
                filters = kwargs['filters']
                if isinstance(filters, dict) and filters:
                    screen_name = ', '.join(f'{k}={v}' for k, v in list(filters.items())[:3])
            
            if screen_name:
                properties['$screen_name'] = screen_name
            
            # Capture event (async, non-blocking)
            posthog.capture(
                distinct_id=distinct_id,
                event=tool_name,
                properties=properties
            )
            
            debug_print(f"[posthog] Captured: {tool_name} for {distinct_id}")
        except Exception as e:
            # Never crash on analytics errors
            debug_print(f"[posthog] Failed to capture event: {e}")
        
        return result
    
    return wrapper


def load_bootstrap() -> str:
    """Load bundled bootstrap.md content.
    
    Returns:
        Bootstrap content as string, or empty string if file missing/unreadable.
    """
    try:
        # Python 3.10+ compatible resource loading
        ref = pkg_resources.files('ims_mcp.resources').joinpath('bootstrap.md')
        with ref.open('r', encoding='utf-8') as f:
            content = f.read()
            debug_print(f"[ims-mcp] v{__version__}: Loaded bootstrap.md ({len(content)} bytes)")
            return content
    except FileNotFoundError:
        debug_print("[ims-mcp] Warning: bootstrap.md not found in package")
        return ""
    except Exception as e:
        debug_print(f"[ims-mcp] Warning: Could not load bootstrap.md: {e}")
        return ""


# Load bootstrap content once at module level (cached)
BOOTSTRAP_CONTENT = load_bootstrap()


def get_authenticated_client() -> R2RClient:
    """Get or create an authenticated R2R client.
    
    This function handles authentication using either:
    1. API key (R2R_API_KEY) - preferred method
    2. Email/password (R2R_EMAIL + R2R_PASSWORD) - fallback method
    
    Returns:
        Authenticated R2RClient instance
    """
    global _authenticated_client
    
    # If client already exists and is authenticated, reuse it
    if _authenticated_client is not None:
        return _authenticated_client
    
    # Log configuration on first client creation
    base_url = os.getenv('R2R_API_BASE') or os.getenv('R2R_BASE_URL') or 'http://localhost:7272'
    collection = os.getenv('R2R_COLLECTION', 'default')
    api_key = os.getenv('R2R_API_KEY', '')
    email = os.getenv("R2R_EMAIL")
    password = os.getenv("R2R_PASSWORD")
    
    debug_print(f"[ims-mcp v{__version__}]")
    debug_print(f"  server={base_url}")
    debug_print(f"  collection={collection}")
    debug_print(f"  api_key={api_key[:3] + '...' if api_key else 'none'}")
    debug_print(f"  email={email if email else 'none'}")
    debug_print(f"  password={password[:3] + '...' if password else 'none'}")
    
    # Create new client
    client = R2RClient()
    
    # Check for email/password authentication
    email = os.getenv("R2R_EMAIL")
    password = os.getenv("R2R_PASSWORD")
    
    if email and password:
        try:
            # Login - R2RClient automatically handles token internally
            client.users.login(email=email, password=password)
            debug_print(f"[ims-mcp] Login successful")
        except Exception as e:
            debug_print(f"[ims-mcp] Login failed: {e}")
            # If login fails, continue without authentication (might work for local servers)
            pass
    
    # Cache the client for reuse
    _authenticated_client = client
    return client


def retry_on_auth_error(func):
    """Decorator to handle token expiry and automatically re-authenticate.
    
    When R2R server restarts, all authentication tokens are invalidated.
    This decorator catches 401/403 errors, invalidates the cached client,
    re-authenticates, and retries the request once.
    """
    def invalidate_client():
        """Invalidate cached client to force re-authentication."""
        global _authenticated_client
        _authenticated_client = None
    
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        try:
            return await func(*args, **kwargs)
        except R2RException as e:
            # Check if this is an authentication error (token expired)
            if hasattr(e, 'status_code') and e.status_code in [401, 403]:
                debug_print(f"[ims-mcp] Token expired, re-authenticating...")
                invalidate_client()
                # Retry once with fresh authentication
                return await func(*args, **kwargs)
            # Re-raise non-auth errors
            raise
    return wrapper


def id_to_shorthand(id: str) -> str:
    """Convert a full ID to shortened version for display."""
    return str(id)[:7]


def format_search_results_for_llm(results) -> str:
    """
    Format R2R search results for LLM consumption.
    
    Formats vector search, graph search, web search, and document search results
    into a readable text format with source IDs and relevant metadata.
    """
    lines = []

    # 1) Chunk search
    if results.chunk_search_results:
        lines.append("Vector Search Results:")
        for c in results.chunk_search_results:
            lines.append(f"Source ID [{id_to_shorthand(c.id)}]:")
            lines.append(c.text or "")  # or c.text[:200] to truncate

    # 2) Graph search
    if results.graph_search_results:
        lines.append("Graph Search Results:")
        for g in results.graph_search_results:
            lines.append(f"Source ID [{id_to_shorthand(g.id)}]:")
            if hasattr(g.content, "summary"):
                lines.append(f"Community Name: {g.content.name}")
                lines.append(f"ID: {g.content.id}")
                lines.append(f"Summary: {g.content.summary}")
            elif hasattr(g.content, "name") and hasattr(
                g.content, "description"
            ):
                lines.append(f"Entity Name: {g.content.name}")
                lines.append(f"Description: {g.content.description}")
            elif (
                hasattr(g.content, "subject")
                and hasattr(g.content, "predicate")
                and hasattr(g.content, "object")
            ):
                lines.append(
                    f"Relationship: {g.content.subject}-{g.content.predicate}-{g.content.object}"
                )

    # 3) Web search
    if results.web_search_results:
        lines.append("Web Search Results:")
        for w in results.web_search_results:
            lines.append(f"Source ID [{id_to_shorthand(w.id)}]:")
            lines.append(f"Title: {w.title}")
            lines.append(f"Link: {w.link}")
            lines.append(f"Snippet: {w.snippet}")

    # 4) Local context docs
    if results.document_search_results:
        lines.append("Local Context Documents:")
        for doc_result in results.document_search_results:
            doc_title = doc_result.title or "Untitled Document"
            doc_id = doc_result.id
            summary = doc_result.summary

            lines.append(f"Full Document ID: {doc_id}")
            lines.append(f"Shortened Document ID: {id_to_shorthand(doc_id)}")
            lines.append(f"Document Title: {doc_title}")
            if summary:
                lines.append(f"Summary: {summary}")

            if doc_result.chunks:
                # Then each chunk inside:
                for chunk in doc_result.chunks:
                    lines.append(
                        f"\nChunk ID {id_to_shorthand(chunk['id'])}:\n{chunk['text']}"
                    )

    result = "\n".join(lines)
    return result


# Create a FastMCP server
try:
    from mcp.server.fastmcp import FastMCP, Context

    mcp = FastMCP(
        name="Rosetta",
        instructions=BOOTSTRAP_CONTENT
    )
except Exception as e:
    raise ImportError(
        "MCP is not installed. Please run `pip install mcp`"
    ) from e


# Search tool with filtering support
@mcp.tool()
@retry_on_auth_error
@track_tool_call
async def search(
    query: str,
    filters: dict | None = None,
    limit: float | None = None,  # Use float to accept JSON "number" type, convert to int internally
    use_semantic_search: bool | None = None,
    use_fulltext_search: bool | None = None,
    ctx: Context = None,
) -> str:
    """
    Performs a search with optional filtering and configuration

    Args:
        query: The search query
        filters: Metadata filters (e.g., {"tags": {"$in": ["agents"]}})
        limit: Maximum number of results (server default if not specified)
        use_semantic_search: Enable semantic search (server default if not specified)
        use_fulltext_search: Enable fulltext search (server default if not specified)

    Returns:
        Formatted search results from the knowledge base
    """
    client = get_authenticated_client()

    # Only build search_settings if user provided any parameters
    # This preserves original behavior: search("query") → search(query=query) with NO search_settings
    kwargs = {"query": query}

    if any(
        param is not None
        for param in [filters, limit, use_semantic_search, use_fulltext_search]
    ):
        search_settings = {}

        if filters is not None:
            search_settings["filters"] = filters
        if limit is not None:
            search_settings["limit"] = int(limit)  # Convert to int for R2R API
        if use_semantic_search is not None:
            search_settings["use_semantic_search"] = use_semantic_search
        if use_fulltext_search is not None:
            search_settings["use_fulltext_search"] = use_fulltext_search

        kwargs["search_settings"] = search_settings

    search_response = client.retrieval.search(**kwargs)
    return format_search_results_for_llm(search_response.results)


# RAG query tool with filtering and generation config
@mcp.tool()
@retry_on_auth_error
@track_tool_call
async def rag(
    query: str,
    filters: dict | None = None,
    limit: float | None = None,  # Use float to accept JSON "number" type, convert to int internally
    model: str | None = None,
    temperature: float | None = None,
    max_tokens: float | None = None,  # Use float to accept JSON "number" type, convert to int internally
    ctx: Context = None,
) -> str:
    """
    Perform RAG query with optional filtering and generation config

    Args:
        query: The question to answer
        filters: Metadata filters (e.g., {"tags": {"$in": ["agents"]}})
        limit: Max search results to use (server default if not specified)
        model: LLM model to use (server default if not specified)
        temperature: Response randomness 0-1 (server default if not specified)
        max_tokens: Max response length (server default if not specified)

    Returns:
        Generated answer from RAG
    """
    client = get_authenticated_client()

    # Only build configs if user provided parameters
    # This preserves original behavior: rag("query") → rag(query=query) with NO configs
    kwargs = {"query": query}

    # Build search_settings if any search params provided
    if any(param is not None for param in [filters, limit]):
        search_settings = {}
        if filters is not None:
            search_settings["filters"] = filters
        if limit is not None:
            search_settings["limit"] = int(limit)  # Convert to int for R2R API
        kwargs["search_settings"] = search_settings

    # Build rag_generation_config if any generation params provided
    if any(param is not None for param in [model, temperature, max_tokens]):
        rag_config = {}
        if model is not None:
            rag_config["model"] = model
        if temperature is not None:
            rag_config["temperature"] = temperature
        if max_tokens is not None:
            rag_config["max_tokens"] = int(max_tokens)  # Convert to int for R2R API
        kwargs["rag_generation_config"] = rag_config

    rag_response = client.retrieval.rag(**kwargs)
    return rag_response.results.generated_answer  # type: ignore


# Document upload tool with upsert semantics
#@mcp.tool() # disabled intentionally to prevent accidental document uploads, and because R2R does not support proper permissions management.
@retry_on_auth_error
@track_tool_call
async def put_document(
    content: str,
    title: str,
    metadata: dict | None = None,
    document_id: str | None = None,
    ctx: Context = None,
) -> str:
    """
    Upload or update a document with upsert semantics

    Args:
        content: The text content of the document
        title: Document title (used for ID generation if document_id not provided)
        metadata: Additional metadata (e.g., {"tags": ["agents"], "domain": "dev"})
        document_id: Optional document ID for explicit upsert

    Returns:
        Status message with document_id and operation type
    """
    client = get_authenticated_client()

    # Build metadata with title
    final_metadata = {"title": title}
    if metadata:
        final_metadata.update(metadata)

    # Generate UUID from title if not provided (enables upsert by title)
    if not document_id:
        # Use UUID5 to create deterministic UUID from title
        # Same title always generates same UUID for upsert semantics
        namespace = uuid.UUID("6ba7b810-9dad-11d1-80b4-00c04fd430c8")  # DNS namespace
        document_id = str(uuid.uuid5(namespace, title))

    # Try to create, if exists then update
    try:
        result = client.documents.create(
            raw_text=content,
            id=document_id,
            metadata=final_metadata,
            run_with_orchestration=True,
        )
        return f"Document created successfully.\nDocument ID: {document_id}\nOperation: CREATE"
    except Exception as e:
        error_msg = str(e)
        # Check if document already exists
        if "already exists" in error_msg.lower():
            # Document exists, delete and recreate for true upsert
            try:
                # Delete existing document
                client.documents.delete(id=document_id)
                
                # Recreate with new content
                result = client.documents.create(
                    raw_text=content,
                    id=document_id,
                    metadata=final_metadata,
                    run_with_orchestration=True,
                )
                return f"Document updated successfully.\nDocument ID: {document_id}\nOperation: UPDATE (delete + recreate)"
            except Exception as update_error:
                return f"Error updating document: {str(update_error)}"
        else:
            # Different error, re-raise
            return f"Error creating document: {error_msg}"


# List documents tool
@mcp.tool()
@retry_on_auth_error
@track_tool_call
async def list_documents(
    offset: float = 0,  # Use float to accept JSON "number" type, convert to int internally
    limit: float = 100,  # Use float to accept JSON "number" type, convert to int internally
    document_ids: list[str] | None = None,
    compact_view: bool = True,
    tags: list[str] | None = None,
    match_all_tags: bool = False,
    ctx: Context = None,
) -> str:
    """
    List documents in the R2R knowledge base with pagination

    Args:
        offset: Number of documents to skip (default: 0)
        limit: Maximum number of documents to return (default: 100, max: 100)
        document_ids: Optional list of specific document IDs to retrieve
        compact_view: Show only ID and title (default: True - compact view)
        tags: Optional list of tags to filter by (e.g., ["agents", "r1"])
        match_all_tags: If True, document must have ALL tags; if False (default), document must have ANY tag

    Returns:
        Formatted list of documents
    """
    client = get_authenticated_client()

    # Build kwargs for list call
    kwargs = {"offset": int(offset), "limit": min(int(limit), 100)}  # Convert to int, cap at 100
    
    if document_ids:
        kwargs["ids"] = document_ids

    # List documents
    result = client.documents.list(**kwargs)
    
    # Filter by tags if provided
    filtered_results = result.results
    if tags and len(tags) > 0:
        provided_tags = set(tags)
        filtered_results = []
        
        for doc in result.results:
            # Extract tags from document metadata
            doc_tags = set()
            if hasattr(doc, 'metadata') and doc.metadata:
                tags_value = doc.metadata.get('tags')
                if tags_value:
                    if isinstance(tags_value, list):
                        doc_tags = set(tags_value)
                    elif isinstance(tags_value, str):
                        doc_tags = {tags_value}
            
            # Apply filter based on match mode
            if match_all_tags:
                # ALL mode: document must have all provided tags
                if provided_tags.issubset(doc_tags):
                    filtered_results.append(doc)
            else:
                # ANY mode: document must have at least one provided tag
                if len(provided_tags.intersection(doc_tags)) > 0:
                    filtered_results.append(doc)
    
    # Format results for display
    lines = []
    if tags:
        tag_mode = "ALL" if match_all_tags else "ANY"
        lines.append(f"Documents with {tag_mode} tags {tags} (showing {len(filtered_results)} of {result.total_entries} total):\n")
    else:
        lines.append(f"Documents (showing {len(filtered_results)} of {result.total_entries} total):\n")
    
    for doc in filtered_results:
        if compact_view:
            # Compact mode: just ID and title
            lines.append(f"ID: {doc.id} | Title: {doc.title or 'Untitled'}")
        else:
            # Full mode: all details
            lines.append(f"{'='*60}")
            lines.append(f"ID: {doc.id}")
            lines.append(f"Title: {doc.title or 'Untitled'}")
            
            # Show all metadata
            if hasattr(doc, 'metadata') and doc.metadata:
                lines.append("Metadata:")
                for key, value in doc.metadata.items():
                    # Format value based on type
                    if isinstance(value, list):
                        lines.append(f"  {key}: {', '.join(str(v) for v in value)}")
                    elif isinstance(value, dict):
                        lines.append(f"  {key}: {value}")
                    else:
                        lines.append(f"  {key}: {value}")
            
            lines.append(f"Status: {doc.ingestion_status}")
            lines.append(f"Size: {doc.size_in_bytes} bytes")
            lines.append(f"Created: {doc.created_at}")
            lines.append(f"Updated: {doc.updated_at}")
            
            if hasattr(doc, 'summary') and doc.summary:
                lines.append(f"Summary: {doc.summary[:200]}...")
    
    return "\n".join(lines)


# Get document tool
@mcp.tool()
@retry_on_auth_error
@track_tool_call
async def get_document(
    document_id: str | None = None,
    title: str | None = None,
    ctx: Context = None,
) -> str:
    """
    Retrieve a document by ID or title
    
    Args:
        document_id: Document ID to retrieve
        title: Document title to search for (if document_id not provided)
    
    Returns:
        Formatted document details
    """
    client = get_authenticated_client()
    
    # Validate that at least one parameter is provided
    if not document_id and not title:
        return "Error: Must provide either document_id or title"
    
    # If only title provided, search for document by title using metadata filter
    if not document_id and title:
        try:
            # Use search API with title filter for more efficient lookup
            search_result = client.retrieval.search(
                query=title,
                search_settings={
                    "filters": {"title": {"$eq": title}},
                    "limit": 5,
                    "use_fulltext_search": True,  # Use fulltext for exact match
                }
            )
            
            # Extract document IDs from search results
            matching_docs = []
            if hasattr(search_result, 'results') and hasattr(search_result.results, 'chunk_search_results'):
                seen_doc_ids = set()
                for chunk in search_result.results.chunk_search_results:
                    if hasattr(chunk, 'metadata') and chunk.metadata:
                        doc_id = chunk.metadata.get('document_id')
                        doc_title = chunk.metadata.get('title')
                        if doc_id and doc_id not in seen_doc_ids:
                            seen_doc_ids.add(doc_id)
                            matching_docs.append((doc_id, doc_title))
            
            # Fallback: Use list API if search didn't work
            if not matching_docs:
                list_result = client.documents.list(limit=100)
                for doc in list_result.results:
                    if doc.title and doc.title.lower() == title.lower():
                        matching_docs.append((doc.id, doc.title))
            
            if not matching_docs:
                return f"Error: No document found with title '{title}'"
            
            if len(matching_docs) > 1:
                lines = [f"Warning: Found {len(matching_docs)} documents with title '{title}':"]
                for doc_id, doc_title in matching_docs:
                    lines.append(f"  - ID: {doc_id}")
                lines.append("\nPlease use document_id to retrieve a specific document.")
                return "\n".join(lines)
            
            # Found exactly one match
            document_id = matching_docs[0][0]
        except R2RException as e:
            # Re-raise authentication errors so the decorator can handle them
            if hasattr(e, 'status_code') and e.status_code in [401, 403]:
                raise
            # For other R2RExceptions, return error message
            return f"Error searching for document by title: {str(e)}"
        except Exception as e:
            return f"Error searching for document by title: {str(e)}"
    
    # Now we have document_id, download the original file
    try:
        # Download the original document content
        file_content = client.documents.download(id=document_id)
        
        # Read and decode the content
        content_bytes = file_content.read()
        content_text = content_bytes.decode('utf-8')
        
        # Build output with just document ID and content
        output_lines = [f"DOCUMENT ID: {document_id}", content_text]
        
        return "\n".join(output_lines)
        
    except R2RException as e:
        # Re-raise authentication errors so the decorator can handle them
        if hasattr(e, 'status_code') and e.status_code in [401, 403]:
            raise
        # For other R2RExceptions, return error message
        error_msg = str(e)
        if "not found" in error_msg.lower():
            return f"Error: Document with ID '{document_id}' not found"
        return f"Error downloading document: {error_msg}"
    except Exception as e:
        error_msg = str(e)
        if "not found" in error_msg.lower():
            return f"Error: Document with ID '{document_id}' not found"
        return f"Error downloading document: {error_msg}"


# Delete document tool
#@mcp.tool() # disabled intentionally to prevent accidental document uploads, and because R2R does not support proper permissions management.
@retry_on_auth_error
@track_tool_call
async def delete_document(
    document_id: str,
    ctx: Context = None,
) -> str:
    """
    Delete a document by ID
    
    Args:
        document_id: The unique identifier of the document to delete
    
    Returns:
        Status message confirming deletion or describing error
    """
    client = get_authenticated_client()
    
    try:
        # Delete the document
        client.documents.delete(id=document_id)
        return f"Document deleted successfully.\nDocument ID: {document_id}"
    except Exception as e:
        error_msg = str(e)
        
        # Handle common error cases with user-friendly messages
        if "not found" in error_msg.lower():
            return f"Error: Document with ID '{document_id}' not found"
        elif "permission" in error_msg.lower() or "denied" in error_msg.lower():
            return f"Error: Permission denied to delete document '{document_id}'"
        elif "connection" in error_msg.lower() or "timeout" in error_msg.lower():
            return f"Error: Unable to communicate with R2R server. Please check connection."
        else:
            # Generic error fallback
            return f"Error deleting document: {error_msg}"


def main():
    """Main entry point for console script."""
    mcp.run()


# Run the server if executed directly
if __name__ == "__main__":
    main()

