"""
MCP Tools for REM operations.

Tools are functions that LLMs can call to interact with the REM system.
Each tool is decorated with @mcp.tool() and registered with the FastMCP server.

Design Pattern:
- Tools receive parameters from LLM
- Tools delegate to RemService or ContentService
- Tools return structured results
- Tools handle errors gracefully with informative messages

Available Tools:
- search_rem: Execute REM queries (LOOKUP, FUZZY, SEARCH, SQL, TRAVERSE)
- ask_rem_agent: Natural language to REM query conversion via agent
- ingest_into_rem: Full file ingestion pipeline (read + store + parse + chunk)
- read_resource: Access MCP resources (for Claude Desktop compatibility)
- register_metadata: Register response metadata for SSE MetadataEvent
- list_schema: List all schemas (tables, agents) in the database with row counts
- get_schema: Get detailed schema for a table (columns, types, indexes)
"""

import json
from functools import wraps
from typing import Any, Callable, Literal, cast

from loguru import logger

from ...agentic.context import AgentContext
from ...models.core import (
    FuzzyParameters,
    LookupParameters,
    QueryType,
    RemQuery,
    SearchParameters,
    SQLParameters,
    TraverseParameters,
)
from ...services.postgres import PostgresService
from ...services.rem import RemService
from ...settings import settings


# Service cache for FastAPI lifespan initialization
_service_cache: dict[str, Any] = {}


def init_services(postgres_service: PostgresService, rem_service: RemService):
    """
    Initialize service instances for MCP tools.

    Called during FastAPI lifespan startup.

    Args:
        postgres_service: PostgresService instance
        rem_service: RemService instance
    """
    _service_cache["postgres"] = postgres_service
    _service_cache["rem"] = rem_service
    logger.debug("MCP tools initialized with service instances")


async def get_rem_service() -> RemService:
    """
    Get or create RemService instance (lazy initialization).

    Returns cached instance if available, otherwise creates new one.
    Always ensures PostgreSQL connection is active before returning.
    Thread-safe for async usage.
    """
    # Lazy initialization for in-process/CLI usage
    from ...services.postgres import get_postgres_service

    postgres_service = get_postgres_service()
    if not postgres_service:
        raise RuntimeError("PostgreSQL is disabled. Cannot use REM service.")

    # Always ensure connection is active (handles reconnection if pool was closed)
    await postgres_service.connect()

    if "rem" in _service_cache:
        return cast(RemService, _service_cache["rem"])

    rem_service = RemService(postgres_service=postgres_service)

    _service_cache["postgres"] = postgres_service
    _service_cache["rem"] = rem_service

    logger.debug("MCP tools: lazy initialized services")
    return rem_service


def mcp_tool_error_handler(func: Callable) -> Callable:
    """
    Decorator for consistent MCP tool error handling.

    Wraps tool functions to:
    - Log errors with full context
    - Return standardized error responses
    - Prevent exceptions from bubbling to LLM

    Usage:
        @mcp_tool_error_handler
        async def my_tool(...) -> dict[str, Any]:
            # Pure business logic - no try/except needed
            result = await service.do_work()
            return {"data": result}

    Returns:
        {"status": "success", **result} on success
        {"status": "error", "error": str(e)} on failure
    """
    @wraps(func)
    async def wrapper(*args, **kwargs) -> dict[str, Any]:
        try:
            result = await func(*args, **kwargs)
            # If result already has status, return as-is
            if isinstance(result, dict) and "status" in result:
                return result
            # Otherwise wrap in success response
            return {"status": "success", **result}
        except Exception as e:
            # Use %s format to avoid issues with curly braces in error messages
            logger.opt(exception=True).error("{} failed: {}", func.__name__, str(e))
            return {
                "status": "error",
                "error": str(e),
                "tool": func.__name__,
            }
    return wrapper


@mcp_tool_error_handler
async def search_rem(
    query: str,
    limit: int = 20,
) -> dict[str, Any]:
    """
    Execute a REM query using the REM query dialect.

    **REM Query Syntax:**

    LOOKUP <entity_key>
        Find entity by exact name/key. Searches across all tables.
        Example: LOOKUP phq-9-procedure
        Example: LOOKUP sertraline

    SEARCH <text> IN <table>
        Semantic vector search within a specific table.
        Tables: 'ontologies' (clinical knowledge, procedures, drugs, DSM criteria)
                'resources' (documents, files, user content)
        Example: SEARCH depression IN ontologies
        Example: SEARCH Module F IN ontologies

    FUZZY <text>
        Fuzzy text matching for partial matches and typos.
        Example: FUZZY setraline

    TRAVERSE <start_entity>
        Graph traversal from a starting entity.
        Example: TRAVERSE sarah-chen

    Args:
        query: REM query string (e.g., "LOOKUP phq-9-procedure", "SEARCH depression IN ontologies")
        limit: Maximum results to return (default: 20)

    Returns:
        Dict with query results and metadata. If no results found, includes
        'suggestions' with alternative search strategies.

    Examples:
        search_rem("LOOKUP phq-9-procedure")
        search_rem("SEARCH depression IN ontologies")
        search_rem("SEARCH anxiety treatment IN ontologies", limit=10)
        search_rem("FUZZY setraline")
    """
    # Get RemService instance (lazy initialization)
    rem_service = await get_rem_service()

    # Get user_id from parent context (set by tool wrapper or parent agent)
    from ...agentic.context import get_current_context
    parent_context = get_current_context()
    if parent_context is not None and parent_context.user_id:
        user_id = parent_context.user_id
        logger.debug(f"search_rem: using user_id from context: {user_id}")
    else:
        user_id = AgentContext.get_user_id_or_default(None, source="search_rem")
        logger.debug(f"search_rem: no context, falling back to default: {user_id}")

    # Parse the REM query string
    if not query or not query.strip():
        return {
            "status": "error",
            "error": "Empty query. Use REM syntax: LOOKUP <key>, SEARCH <text> IN <table>, FUZZY <text>, or TRAVERSE <entity>",
        }

    query = query.strip()
    parts = query.split(None, 1)  # Split on first whitespace

    if len(parts) < 2:
        return {
            "status": "error",
            "error": f"Invalid query format: '{query}'. Expected: LOOKUP <key>, SEARCH <text> IN <table>, FUZZY <text>, or TRAVERSE <entity>",
        }

    query_type = parts[0].upper()
    remainder = parts[1].strip()

    # Build RemQuery based on query_type
    if query_type == "LOOKUP":
        if not remainder:
            return {
                "status": "error",
                "error": "LOOKUP requires an entity key. Example: LOOKUP phq-9-procedure",
            }

        rem_query = RemQuery(
            query_type=QueryType.LOOKUP,
            parameters=LookupParameters(
                key=remainder,
                user_id=user_id,
            ),
            user_id=user_id,
        )
        table = None  # LOOKUP searches all tables

    elif query_type == "SEARCH":
        # Parse "text IN table" format
        if " IN " in remainder.upper():
            # Find the last " IN " to handle cases like "SEARCH pain IN back IN ontologies"
            in_pos = remainder.upper().rfind(" IN ")
            search_text = remainder[:in_pos].strip()
            table = remainder[in_pos + 4:].strip().lower()
        else:
            return {
                "status": "error",
                "error": f"SEARCH requires table: SEARCH <text> IN <table>. "
                "Use 'ontologies' for clinical knowledge or 'resources' for documents. "
                f"Example: SEARCH {remainder} IN ontologies",
            }

        if not search_text:
            return {
                "status": "error",
                "error": "SEARCH requires search text. Example: SEARCH depression IN ontologies",
            }

        rem_query = RemQuery(
            query_type=QueryType.SEARCH,
            parameters=SearchParameters(
                query_text=search_text,
                table_name=table,
                limit=limit,
            ),
            user_id=user_id,
        )

    elif query_type == "FUZZY":
        if not remainder:
            return {
                "status": "error",
                "error": "FUZZY requires search text. Example: FUZZY setraline",
            }

        rem_query = RemQuery(
            query_type=QueryType.FUZZY,
            parameters=FuzzyParameters(
                query_text=remainder,
                threshold=0.3,  # pg_trgm similarity - 0.3 is reasonable for typo correction
                limit=limit,
            ),
            user_id=user_id,
        )
        table = None

    elif query_type == "TRAVERSE":
        if not remainder:
            return {
                "status": "error",
                "error": "TRAVERSE requires a starting entity. Example: TRAVERSE sarah-chen",
            }

        rem_query = RemQuery(
            query_type=QueryType.TRAVERSE,
            parameters=TraverseParameters(
                initial_query=remainder,
                edge_types=[],
                max_depth=1,
            ),
            user_id=user_id,
        )
        table = None

    else:
        return {
            "status": "error",
            "error": f"Unknown query type: '{query_type}'. Valid types: LOOKUP, SEARCH, FUZZY, TRAVERSE. "
            "Examples: LOOKUP phq-9-procedure, SEARCH depression IN ontologies",
        }

    # Execute query (errors handled by decorator)
    logger.info(f"Executing REM query: {query_type} for user {user_id}")
    result = await rem_service.execute_query(rem_query)

    logger.info(f"Query completed successfully: {query_type}")

    # Provide helpful guidance when no results found
    response: dict[str, Any] = {
        "query_type": query_type,
        "results": result,
    }

    # Check if results are empty - handle both list and dict result formats
    is_empty = False
    if not result:
        is_empty = True
    elif isinstance(result, list) and len(result) == 0:
        is_empty = True
    elif isinstance(result, dict):
        # RemService returns dict with 'results' key containing actual matches
        inner_results = result.get("results", [])
        count = result.get("count", len(inner_results) if isinstance(inner_results, list) else 0)
        is_empty = count == 0 or (isinstance(inner_results, list) and len(inner_results) == 0)

    if is_empty:
        # Build helpful suggestions based on query type
        suggestions = []

        if query_type in ("LOOKUP", "FUZZY"):
            suggestions.append(
                "LOOKUP/FUZZY searches across ALL tables. If you expected results, "
                "verify the entity name is spelled correctly."
            )

        if query_type == "SEARCH":
            if table == "resources":
                suggestions.append(
                    "No results in 'resources' table. Try: SEARCH <text> IN ontologies - "
                    "clinical procedures, drug info, and diagnostic criteria are stored there."
                )
            elif table == "ontologies":
                suggestions.append(
                    "No results in 'ontologies' table. Try: SEARCH <text> IN resources - "
                    "for user-uploaded documents and general content."
                )
            else:
                suggestions.append(
                    "Try: SEARCH <text> IN ontologies (clinical knowledge, procedures, drugs) "
                    "or SEARCH <text> IN resources (documents, files)."
                )

        # Always suggest both tables if no specific table guidance given
        if not suggestions:
            suggestions.append(
                "No results found. Try: SEARCH <text> IN ontologies (clinical procedures, drugs) "
                "or SEARCH <text> IN resources (documents, files)."
            )

        response["suggestions"] = suggestions
        response["hint"] = "0 results returned. See 'suggestions' for alternative search strategies."

    return response


@mcp_tool_error_handler
async def ask_rem_agent(
    query: str,
    agent_schema: str = "ask_rem",
    agent_version: str | None = None,
    user_id: str | None = None,
) -> dict[str, Any]:
    """
    Ask REM using natural language via agent-driven query conversion.

    This tool converts natural language questions into optimized REM queries
    using an agent that understands the REM query language and schema.

    The agent can perform multi-turn reasoning and iterated retrieval:
    1. Initial exploration (LOOKUP/FUZZY to find entities)
    2. Semantic search (SEARCH for related content)
    3. Graph traversal (TRAVERSE to explore relationships)
    4. Synthesis (combine results into final answer)

    Args:
        query: Natural language question or task
        agent_schema: Agent schema name (default: "ask_rem")
        agent_version: Optional agent version (default: latest)
        user_id: Optional user identifier (defaults to authenticated user or "default")

    Returns:
        Dict with:
        - status: "success" or "error"
        - response: Agent's natural language response
        - query_output: Structured query results (if available)
        - queries_executed: List of REM queries executed
        - metadata: Agent execution metadata

    Examples:
        # Simple question (uses authenticated user context)
        ask_rem_agent(
            query="Who is Sarah Chen?"
        )

        # Complex multi-step question
        ask_rem_agent(
            query="What are the key findings from last week's sprint retrospective?"
        )

        # Graph exploration
        ask_rem_agent(
            query="Show me Sarah's reporting chain and their recent projects"
        )
    """
    from ...agentic import create_agent
    from ...agentic.context import get_current_context
    from ...utils.schema_loader import load_agent_schema

    # Get parent context for multi-agent support
    # This enables context propagation from parent agent to child agent
    parent_context = get_current_context()

    # Build child context: inherit from parent if available, otherwise use defaults
    if parent_context is not None:
        # Inherit user_id, tenant_id, session_id, is_eval from parent
        # Allow explicit user_id override if provided
        effective_user_id = user_id or parent_context.user_id
        context = parent_context.child_context(agent_schema_uri=agent_schema)
        if user_id is not None:
            # Override user_id if explicitly provided
            context = AgentContext(
                user_id=user_id,
                tenant_id=parent_context.tenant_id,
                session_id=parent_context.session_id,
                default_model=parent_context.default_model,
                agent_schema_uri=agent_schema,
                is_eval=parent_context.is_eval,
            )
        logger.debug(
            f"ask_rem_agent inheriting context from parent: "
            f"user_id={context.user_id}, session_id={context.session_id}"
        )
    else:
        # No parent context - create fresh context (backwards compatible)
        effective_user_id = AgentContext.get_user_id_or_default(
            user_id, source="ask_rem_agent"
        )
        context = AgentContext(
            user_id=effective_user_id,
            tenant_id=effective_user_id or "default",
            default_model=settings.llm.default_model,
            agent_schema_uri=agent_schema,
        )

    # Load agent schema
    try:
        schema = load_agent_schema(agent_schema)
    except FileNotFoundError:
        return {
            "status": "error",
            "error": f"Agent schema not found: {agent_schema}",
        }

    # Create agent
    agent_runtime = await create_agent(
        context=context,
        agent_schema_override=schema,
    )

    # Run agent (errors handled by decorator)
    logger.debug(f"Running ask_rem agent for query: {query[:100]}...")
    result = await agent_runtime.run(query)

    # Extract output
    from rem.agentic.serialization import serialize_agent_result
    query_output = serialize_agent_result(result.output)

    logger.debug("Agent execution completed successfully")

    return {
        "response": str(result.output),
        "query_output": query_output,
        "natural_query": query,
    }


@mcp_tool_error_handler
async def ingest_into_rem(
    file_uri: str,
    category: str | None = None,
    tags: list[str] | None = None,
    is_local_server: bool = False,
    resource_type: str | None = None,
) -> dict[str, Any]:
    """
    Ingest file into REM, creating searchable PUBLIC resources and embeddings.

    **IMPORTANT: All ingested data is PUBLIC by default.** This is correct for
    shared knowledge bases (ontologies, procedures, reference data). Private
    user-scoped data requires different handling via the CLI with --make-private.

    This tool provides the complete file ingestion pipeline:
    1. **Read**: File from local/S3/HTTP
    2. **Store**: To internal storage (public namespace)
    3. **Parse**: Extract content, metadata, tables, images
    4. **Chunk**: Semantic chunking for embeddings
    5. **Embed**: Create Resource chunks with vector embeddings

    Supported file types:
    - Documents: PDF, DOCX, TXT, Markdown
    - Code: Python, JavaScript, TypeScript, etc.
    - Data: CSV, JSON, YAML
    - Audio: WAV, MP3 (transcription)

    **Security**: Remote MCP servers cannot read local files. Only local/stdio
    MCP servers can access local filesystem paths.

    Args:
        file_uri: File location (local path, s3:// URI, or http(s):// URL)
        category: Optional category (document, code, audio, etc.)
        tags: Optional tags for file
        is_local_server: True if running as local/stdio MCP server
        resource_type: Optional resource type for storing chunks (case-insensitive).
            Supports flexible naming:
            - "resource", "resources", "Resource" → Resource (default)
            - "domain-resource", "domain_resource", "DomainResource",
              "domain-resources" → DomainResource (curated internal knowledge)

    Returns:
        Dict with:
        - status: "success" or "error"
        - file_id: Created file UUID
        - file_name: Original filename
        - storage_uri: Internal storage URI
        - processing_status: "completed" or "failed"
        - resources_created: Number of Resource chunks created
        - content: Parsed file content (markdown format) if completed
        - message: Human-readable status message

    Examples:
        # Ingest local file (local server only)
        ingest_into_rem(
            file_uri="/Users/me/procedure.pdf",
            category="medical",
            is_local_server=True
        )

        # Ingest from S3
        ingest_into_rem(
            file_uri="s3://bucket/docs/report.pdf"
        )

        # Ingest from HTTP
        ingest_into_rem(
            file_uri="https://example.com/whitepaper.pdf",
            tags=["research", "whitepaper"]
        )

        # Ingest as curated domain knowledge
        ingest_into_rem(
            file_uri="s3://bucket/internal/procedures.pdf",
            resource_type="domain-resource",
            category="procedures"
        )
    """
    from ...services.content import ContentService

    # Data is PUBLIC by default (user_id=None)
    # Private user-scoped data requires CLI with --make-private flag

    # Delegate to ContentService for centralized ingestion (errors handled by decorator)
    content_service = ContentService()
    result = await content_service.ingest_file(
        file_uri=file_uri,
        user_id=None,  # PUBLIC - all ingested data is shared/public
        category=category,
        tags=tags,
        is_local_server=is_local_server,
        resource_type=resource_type,
    )

    logger.debug(
        f"MCP ingestion complete: {result['file_name']} "
        f"(status: {result['processing_status']}, "
        f"resources: {result['resources_created']})"
    )

    return result


@mcp_tool_error_handler
async def read_resource(uri: str) -> dict[str, Any]:
    """
    Read an MCP resource by URI.

    This tool provides automatic access to MCP resources in Claude Desktop.
    Resources contain authoritative, up-to-date reference data.

    **IMPORTANT**: This tool enables Claude Desktop to automatically access
    resources based on query relevance. While FastMCP correctly exposes resources
    via standard MCP resource endpoints, Claude Desktop currently requires manual
    resource attachment. This tool bridges that gap by exposing resource access
    as a tool, which Claude Desktop WILL automatically invoke.

    **Available Resources:**

    Agent Schemas:
    • rem://agents - List all available agent schemas
    • rem://agents/{agent_name} - Get specific agent schema

    Documentation:
    • rem://schema/entities - Entity schemas (Resource, Message, User, File, Moment)
    • rem://schema/query-types - REM query type documentation

    System Status:
    • rem://status - System health and statistics

    Args:
        uri: Resource URI (e.g., "rem://agents", "rem://agents/ask_rem")

    Returns:
        Dict with:
        - status: "success" or "error"
        - uri: Original URI
        - data: Resource data (format depends on resource type)

    Examples:
        # List all agents
        read_resource(uri="rem://agents")

        # Get specific agent
        read_resource(uri="rem://agents/ask_rem")

        # Check system status
        read_resource(uri="rem://status")
    """
    logger.debug(f"Reading resource: {uri}")

    # Import here to avoid circular dependency
    from .resources import load_resource

    # Load resource using the existing resource handler (errors handled by decorator)
    result = await load_resource(uri)

    logger.debug(f"Resource loaded successfully: {uri}")

    # If result is already a dict, return it
    if isinstance(result, dict):
        return {
            "uri": uri,
            "data": result,
        }

    # If result is a string (JSON), parse it
    import json

    try:
        data = json.loads(result)
        return {
            "uri": uri,
            "data": data,
        }
    except json.JSONDecodeError:
        # Return as plain text if not JSON
        return {
            "uri": uri,
            "data": {"content": result},
        }


async def register_metadata(
    confidence: float | None = None,
    references: list[str] | None = None,
    sources: list[str] | None = None,
    flags: list[str] | None = None,
    # Session naming
    session_name: str | None = None,
    # Risk assessment fields (used by specialized agents)
    risk_level: str | None = None,
    risk_score: int | None = None,
    risk_reasoning: str | None = None,
    recommended_action: str | None = None,
    # Generic extension - any additional key-value pairs
    extra: dict[str, Any] | None = None,
    # Agent schema (auto-populated from context if not provided)
    agent_schema: str | None = None,
) -> dict[str, Any]:
    """
    Register response metadata to be emitted as an SSE MetadataEvent.

    Call this tool BEFORE generating your final response to provide structured
    metadata that will be sent to the client alongside your natural language output.
    This allows you to stream conversational responses while still providing
    machine-readable confidence scores, references, and other metadata.

    **Design Pattern**: Agents can call this once before their final response to
    register metadata that the streaming layer will emit as a MetadataEvent.
    This decouples structured metadata from the response format.

    Args:
        confidence: Confidence score (0.0-1.0) for the response quality.
            - 0.9-1.0: High confidence, answer is well-supported
            - 0.7-0.9: Medium confidence, some uncertainty
            - 0.5-0.7: Low confidence, significant gaps
            - <0.5: Very uncertain, may need clarification
        references: List of reference identifiers (file paths, document IDs,
            entity labels) that support the response.
        sources: List of source descriptions (e.g., "REM database",
            "search results", "user context").
        flags: Optional flags for the response (e.g., "needs_review",
            "uncertain", "incomplete", "crisis_alert").

        session_name: Short 1-3 phrase name describing the session topic.
            Used by the UI to label conversations in the sidebar.
            Examples: "Prescription Drug Questions", "AWS Setup Help",
            "Python Code Review", "Travel Planning".

        risk_level: Risk level indicator (e.g., "green", "orange", "red").
            Used by mental health agents for C-SSRS style assessment.
        risk_score: Numeric risk score (e.g., 0-6 for C-SSRS).
        risk_reasoning: Brief explanation of risk assessment.
        recommended_action: Suggested next steps based on assessment.

        extra: Dict of arbitrary additional metadata. Use this for any
            domain-specific fields not covered by the standard parameters.
            Example: {"topics_detected": ["anxiety", "sleep"], "session_count": 5}
        agent_schema: Optional agent schema name. If not provided, automatically
            populated from the current agent context (for multi-agent tracing).

    Returns:
        Dict with:
        - status: "success"
        - _metadata_event: True (marker for streaming layer)
        - All provided fields merged into response

    Examples:
        # High confidence answer with references
        register_metadata(
            confidence=0.95,
            references=["sarah-chen", "q3-report-2024"],
            sources=["REM database lookup"]
        )

        # Risk assessment example
        register_metadata(
            confidence=0.9,
            risk_level="green",
            risk_score=0,
            risk_reasoning="No risk indicators detected in message",
            sources=["mental_health_resources"]
        )

        # Orange risk with recommended action
        register_metadata(
            risk_level="orange",
            risk_score=2,
            risk_reasoning="Passive ideation detected - 'feeling hopeless'",
            recommended_action="Schedule care team check-in within 24-48 hours",
            flags=["care_team_alert"]
        )

        # Custom domain-specific metadata
        register_metadata(
            confidence=0.8,
            extra={
                "topics_detected": ["medication", "side_effects"],
                "drug_mentioned": "sertraline",
                "sentiment": "concerned"
            }
        )
    """
    # Auto-populate agent_schema from context if not provided
    if agent_schema is None:
        from ...agentic.context import get_current_context
        current_context = get_current_context()
        if current_context and current_context.agent_schema_uri:
            agent_schema = current_context.agent_schema_uri

    logger.debug(
        f"Registering metadata: confidence={confidence}, "
        f"risk_level={risk_level}, refs={len(references or [])}, "
        f"sources={len(sources or [])}, agent_schema={agent_schema}"
    )

    result = {
        "status": "success",
        "_metadata_event": True,  # Marker for streaming layer
        "confidence": confidence,
        "references": references,
        "sources": sources,
        "flags": flags,
        "agent_schema": agent_schema,  # Include agent schema for tracing
    }

    # Add session name if provided
    if session_name is not None:
        result["session_name"] = session_name

    # Add risk assessment fields if provided
    if risk_level is not None:
        result["risk_level"] = risk_level
    if risk_score is not None:
        result["risk_score"] = risk_score
    if risk_reasoning is not None:
        result["risk_reasoning"] = risk_reasoning
    if recommended_action is not None:
        result["recommended_action"] = recommended_action

    # Merge any extra fields
    if extra:
        result["extra"] = extra

    return result


@mcp_tool_error_handler
async def list_schema(
    include_system: bool = False,
    user_id: str | None = None,
) -> dict[str, Any]:
    """
    List all schemas (tables) in the REM database.

    Returns metadata about all available tables including their names,
    row counts, and descriptions. Use this to discover what data is
    available before constructing queries.

    Args:
        include_system: If True, include PostgreSQL system tables (pg_*, information_schema).
                       Default False shows only REM application tables.
        user_id: Optional user identifier (defaults to authenticated user or "default")

    Returns:
        Dict with:
        - status: "success" or "error"
        - tables: List of table metadata dicts with:
            - name: Table name
            - schema: Schema name (usually "public")
            - estimated_rows: Approximate row count
            - description: Table comment if available

    Examples:
        # List all REM schemas
        list_schema()

        # Include system tables
        list_schema(include_system=True)
    """
    rem_service = await get_rem_service()
    user_id = AgentContext.get_user_id_or_default(user_id, source="list_schema")

    # Query information_schema for tables
    schema_filter = ""
    if not include_system:
        schema_filter = """
            AND table_schema = 'public'
            AND table_name NOT LIKE 'pg_%'
            AND table_name NOT LIKE '_pg_%'
        """

    query = f"""
        SELECT
            t.table_schema,
            t.table_name,
            pg_catalog.obj_description(
                (quote_ident(t.table_schema) || '.' || quote_ident(t.table_name))::regclass,
                'pg_class'
            ) as description,
            (
                SELECT reltuples::bigint
                FROM pg_class c
                JOIN pg_namespace n ON n.oid = c.relnamespace
                WHERE c.relname = t.table_name
                AND n.nspname = t.table_schema
            ) as estimated_rows
        FROM information_schema.tables t
        WHERE t.table_type = 'BASE TABLE'
        {schema_filter}
        ORDER BY t.table_schema, t.table_name
    """

    # Access postgres service directly from cache
    postgres_service = _service_cache.get("postgres")
    if not postgres_service:
        postgres_service = rem_service._postgres

    rows = await postgres_service.fetch(query)

    tables = []
    for row in rows:
        tables.append({
            "name": row["table_name"],
            "schema": row["table_schema"],
            "estimated_rows": int(row["estimated_rows"]) if row["estimated_rows"] else 0,
            "description": row["description"],
        })

    logger.info(f"Listed {len(tables)} schemas for user {user_id}")

    return {
        "tables": tables,
        "count": len(tables),
    }


@mcp_tool_error_handler
async def get_schema(
    table_name: str,
    include_indexes: bool = True,
    include_constraints: bool = True,
    columns: list[str] | None = None,
    user_id: str | None = None,
) -> dict[str, Any]:
    """
    Get detailed schema information for a specific table.

    Returns column definitions, data types, constraints, and indexes.
    Use this to understand table structure before writing SQL queries.

    Args:
        table_name: Name of the table to inspect (e.g., "resources", "moments")
        include_indexes: Include index information (default True)
        include_constraints: Include constraint information (default True)
        columns: Optional list of specific columns to return. If None, returns all columns.
        user_id: Optional user identifier (defaults to authenticated user or "default")

    Returns:
        Dict with:
        - status: "success" or "error"
        - table_name: Name of the table
        - columns: List of column definitions with:
            - name: Column name
            - type: PostgreSQL data type
            - nullable: Whether NULL is allowed
            - default: Default value if any
            - description: Column comment if available
        - indexes: List of indexes (if include_indexes=True)
        - constraints: List of constraints (if include_constraints=True)
        - primary_key: Primary key column(s)

    Examples:
        # Get full schema for resources table
        get_schema(table_name="resources")

        # Get only specific columns
        get_schema(
            table_name="resources",
            columns=["id", "name", "created_at"]
        )

        # Get schema without indexes
        get_schema(
            table_name="moments",
            include_indexes=False
        )
    """
    rem_service = await get_rem_service()
    user_id = AgentContext.get_user_id_or_default(user_id, source="get_schema")

    # Access postgres service
    postgres_service = _service_cache.get("postgres")
    if not postgres_service:
        postgres_service = rem_service._postgres

    # Verify table exists
    exists_query = """
        SELECT EXISTS (
            SELECT 1 FROM information_schema.tables
            WHERE table_schema = 'public' AND table_name = $1
        )
    """
    exists = await postgres_service.fetchval(exists_query, table_name)
    if not exists:
        return {
            "status": "error",
            "error": f"Table '{table_name}' not found in public schema",
        }

    # Get columns
    columns_filter = ""
    if columns:
        placeholders = ", ".join(f"${i+2}" for i in range(len(columns)))
        columns_filter = f"AND column_name IN ({placeholders})"

    columns_query = f"""
        SELECT
            c.column_name,
            c.data_type,
            c.udt_name,
            c.is_nullable,
            c.column_default,
            c.character_maximum_length,
            c.numeric_precision,
            pg_catalog.col_description(
                (quote_ident(c.table_schema) || '.' || quote_ident(c.table_name))::regclass,
                c.ordinal_position
            ) as description
        FROM information_schema.columns c
        WHERE c.table_schema = 'public'
        AND c.table_name = $1
        {columns_filter}
        ORDER BY c.ordinal_position
    """

    params = [table_name]
    if columns:
        params.extend(columns)

    column_rows = await postgres_service.fetch(columns_query, *params)

    column_defs = []
    for row in column_rows:
        # Build a more readable type string
        data_type = row["data_type"]
        if row["character_maximum_length"]:
            data_type = f"{data_type}({row['character_maximum_length']})"
        elif row["udt_name"] in ("int4", "int8", "float4", "float8"):
            # Use common type names
            type_map = {"int4": "integer", "int8": "bigint", "float4": "real", "float8": "double precision"}
            data_type = type_map.get(row["udt_name"], data_type)
        elif row["udt_name"] == "vector":
            data_type = "vector"

        column_defs.append({
            "name": row["column_name"],
            "type": data_type,
            "nullable": row["is_nullable"] == "YES",
            "default": row["column_default"],
            "description": row["description"],
        })

    result = {
        "table_name": table_name,
        "columns": column_defs,
        "column_count": len(column_defs),
    }

    # Get primary key
    pk_query = """
        SELECT a.attname as column_name
        FROM pg_index i
        JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
        WHERE i.indrelid = $1::regclass
        AND i.indisprimary
        ORDER BY array_position(i.indkey, a.attnum)
    """
    pk_rows = await postgres_service.fetch(pk_query, table_name)
    result["primary_key"] = [row["column_name"] for row in pk_rows]

    # Get indexes
    if include_indexes:
        indexes_query = """
            SELECT
                i.relname as index_name,
                am.amname as index_type,
                ix.indisunique as is_unique,
                ix.indisprimary as is_primary,
                array_agg(a.attname ORDER BY array_position(ix.indkey, a.attnum)) as columns
            FROM pg_index ix
            JOIN pg_class i ON i.oid = ix.indexrelid
            JOIN pg_class t ON t.oid = ix.indrelid
            JOIN pg_am am ON am.oid = i.relam
            JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey)
            WHERE t.relname = $1
            GROUP BY i.relname, am.amname, ix.indisunique, ix.indisprimary
            ORDER BY i.relname
        """
        index_rows = await postgres_service.fetch(indexes_query, table_name)
        result["indexes"] = [
            {
                "name": row["index_name"],
                "type": row["index_type"],
                "unique": row["is_unique"],
                "primary": row["is_primary"],
                "columns": row["columns"],
            }
            for row in index_rows
        ]

    # Get constraints
    if include_constraints:
        constraints_query = """
            SELECT
                con.conname as constraint_name,
                con.contype as constraint_type,
                array_agg(a.attname ORDER BY array_position(con.conkey, a.attnum)) as columns,
                pg_get_constraintdef(con.oid) as definition
            FROM pg_constraint con
            JOIN pg_class t ON t.oid = con.conrelid
            JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(con.conkey)
            WHERE t.relname = $1
            GROUP BY con.conname, con.contype, con.oid
            ORDER BY con.contype, con.conname
        """
        constraint_rows = await postgres_service.fetch(constraints_query, table_name)

        # Map constraint types to readable names
        type_map = {
            "p": "PRIMARY KEY",
            "u": "UNIQUE",
            "f": "FOREIGN KEY",
            "c": "CHECK",
            "x": "EXCLUSION",
        }

        result["constraints"] = []
        for row in constraint_rows:
            # contype is returned as bytes (char type), decode it
            con_type = row["constraint_type"]
            if isinstance(con_type, bytes):
                con_type = con_type.decode("utf-8")
            result["constraints"].append({
                "name": row["constraint_name"],
                "type": type_map.get(con_type, con_type),
                "columns": row["columns"],
                "definition": row["definition"],
            })

    logger.info(f"Retrieved schema for table '{table_name}' with {len(column_defs)} columns")

    return result


@mcp_tool_error_handler
async def save_agent(
    name: str,
    description: str,
    properties: dict[str, Any] | None = None,
    required: list[str] | None = None,
    tools: list[str] | None = None,
    tags: list[str] | None = None,
    version: str = "1.0.0",
    user_id: str | None = None,
) -> dict[str, Any]:
    """
    Save an agent schema to REM, making it available for use.

    This tool creates or updates an agent definition in the user's schema space.
    The agent becomes immediately available for conversations.

    **Default Tools**: All agents automatically get `search_rem` and `register_metadata`
    tools unless explicitly overridden.

    Args:
        name: Agent name in kebab-case (e.g., "code-reviewer", "sales-assistant").
            Must be unique within the user's schema space.
        description: The agent's system prompt. This is the full instruction set
            that defines the agent's behavior, personality, and capabilities.
            Use markdown formatting for structure.
        properties: Output schema properties as a dict. Each property should have:
            - type: "string", "number", "boolean", "array", "object"
            - description: What this field captures
            Example: {"answer": {"type": "string", "description": "Response to user"}}
            If not provided, defaults to a simple {"answer": {"type": "string"}} schema.
        required: List of required property names. Defaults to ["answer"] if not provided.
        tools: List of tool names the agent can use. Defaults to ["search_rem", "register_metadata"].
        tags: Optional tags for categorizing the agent.
        version: Semantic version string (default: "1.0.0").
        user_id: User identifier for scoping. Uses authenticated user if not provided.

    Returns:
        Dict with:
        - status: "success" or "error"
        - agent_name: Name of the saved agent
        - version: Version saved
        - message: Human-readable status

    Examples:
        # Create a simple agent
        save_agent(
            name="greeting-bot",
            description="You are a friendly greeter. Say hello warmly.",
            properties={"answer": {"type": "string", "description": "Greeting message"}},
            required=["answer"]
        )

        # Create agent with structured output
        save_agent(
            name="sentiment-analyzer",
            description="Analyze sentiment of text provided by the user.",
            properties={
                "answer": {"type": "string", "description": "Analysis explanation"},
                "sentiment": {"type": "string", "enum": ["positive", "negative", "neutral"]},
                "confidence": {"type": "number", "minimum": 0, "maximum": 1}
            },
            required=["answer", "sentiment"],
            tags=["analysis", "nlp"]
        )
    """
    from ...agentic.agents.agent_manager import save_agent as _save_agent

    # Get user_id from context if not provided
    user_id = AgentContext.get_user_id_or_default(user_id, source="save_agent")

    # Delegate to agent_manager
    result = await _save_agent(
        name=name,
        description=description,
        user_id=user_id,
        properties=properties,
        required=required,
        tools=tools,
        tags=tags,
        version=version,
    )

    # Add helpful message for Slack users
    if result.get("status") == "success":
        result["message"] = f"Agent '{name}' saved. Use `/custom-agent {name}` to chat with it."

    return result


# =============================================================================
# Multi-Agent Tools
# =============================================================================


@mcp_tool_error_handler
async def ask_agent(
    agent_name: str,
    input_text: str,
    input_data: dict[str, Any] | None = None,
    user_id: str | None = None,
    timeout_seconds: int = 300,
) -> dict[str, Any]:
    """
    Invoke another agent by name and return its response.

    This tool enables multi-agent orchestration by allowing one agent to call
    another. The child agent inherits the parent's context (user_id, session_id,
    tenant_id, is_eval) for proper scoping and continuity.

    Use Cases:
    - Orchestrator agents that delegate to specialized sub-agents
    - Workflow agents that chain multiple processing steps
    - Ensemble agents that aggregate responses from multiple specialists

    Args:
        agent_name: Name of the agent to invoke. Can be:
            - A user-created agent (saved via save_agent)
            - A system agent (e.g., "ask_rem", "knowledge-query")
        input_text: The user message/query to send to the agent
        input_data: Optional structured input data for the agent
        user_id: Optional user override (defaults to parent's user_id)
        timeout_seconds: Maximum execution time (default: 300s)

    Returns:
        Dict with:
        - status: "success" or "error"
        - output: Agent's structured output (if using output schema)
        - text_response: Agent's text response
        - agent_schema: Name of the invoked agent
        - metadata: Any metadata registered by the agent (confidence, etc.)

    Examples:
        # Simple delegation
        ask_agent(
            agent_name="sentiment-analyzer",
            input_text="I love this product! Best purchase ever."
        )
        # Returns: {"status": "success", "output": {"sentiment": "positive"}, ...}

        # Orchestrator pattern
        ask_agent(
            agent_name="knowledge-query",
            input_text="What are the latest Q3 results?"
        )

        # Chain with structured input
        ask_agent(
            agent_name="summarizer",
            input_text="Summarize this document",
            input_data={"document_id": "doc-123", "max_length": 500}
        )
    """
    import asyncio
    from ...agentic import create_agent
    from ...agentic.context import get_current_context, agent_context_scope, get_event_sink, push_event
    from ...agentic.agents.agent_manager import get_agent
    from ...utils.schema_loader import load_agent_schema

    # Get parent context for inheritance
    parent_context = get_current_context()

    # Determine effective user_id
    if parent_context is not None:
        effective_user_id = user_id or parent_context.user_id
    else:
        effective_user_id = AgentContext.get_user_id_or_default(
            user_id, source="ask_agent"
        )

    # Build child context
    if parent_context is not None:
        child_context = parent_context.child_context(agent_schema_uri=agent_name)
        if user_id is not None:
            # Explicit user_id override
            child_context = AgentContext(
                user_id=user_id,
                tenant_id=parent_context.tenant_id,
                session_id=parent_context.session_id,
                default_model=parent_context.default_model,
                agent_schema_uri=agent_name,
                is_eval=parent_context.is_eval,
            )
        logger.debug(
            f"ask_agent '{agent_name}' inheriting context: "
            f"user_id={child_context.user_id}, session_id={child_context.session_id}"
        )
    else:
        child_context = AgentContext(
            user_id=effective_user_id,
            tenant_id=effective_user_id or "default",
            default_model=settings.llm.default_model,
            agent_schema_uri=agent_name,
        )

    # Try to load agent schema from:
    # 1. Database (user-created or system agents)
    # 2. File system (packaged agents)
    schema = None

    # Try database first
    if effective_user_id:
        schema = await get_agent(agent_name, user_id=effective_user_id)
        if schema:
            logger.debug(f"Loaded agent '{agent_name}' from database")

    # Fall back to file system
    if schema is None:
        try:
            schema = load_agent_schema(agent_name)
            logger.debug(f"Loaded agent '{agent_name}' from file system")
        except FileNotFoundError:
            pass

    if schema is None:
        return {
            "status": "error",
            "error": f"Agent not found: {agent_name}",
            "hint": "Use list_agents to see available agents, or save_agent to create one",
        }

    # Create agent runtime
    agent_runtime = await create_agent(
        context=child_context,
        agent_schema_override=schema,
    )

    # Build prompt with optional input_data
    prompt = input_text
    if input_data:
        prompt = f"{input_text}\n\nInput data: {json.dumps(input_data)}"

    # Load session history for the sub-agent (CRITICAL for multi-turn conversations)
    # Sub-agents need to see the full conversation context, not just the summary
    pydantic_message_history = None
    if child_context.session_id and settings.postgres.enabled:
        try:
            from ...services.session import SessionMessageStore, session_to_pydantic_messages
            from ...agentic.schema import get_system_prompt

            store = SessionMessageStore(user_id=child_context.user_id or "default")
            raw_session_history, _has_partition = await store.load_session_messages(
                session_id=child_context.session_id,
                user_id=child_context.user_id,
                compress_on_load=False,  # Need full data for reconstruction
            )
            if raw_session_history:
                # Extract agent's system prompt from schema
                agent_system_prompt = get_system_prompt(schema) if schema else None
                pydantic_message_history = session_to_pydantic_messages(
                    raw_session_history,
                    system_prompt=agent_system_prompt,
                )
                logger.debug(
                    f"ask_agent '{agent_name}': loaded {len(raw_session_history)} session messages "
                    f"-> {len(pydantic_message_history)} pydantic-ai messages"
                )

                # Audit session history if enabled
                from ...services.session import audit_session_history
                audit_session_history(
                    session_id=child_context.session_id,
                    agent_name=agent_name,
                    prompt=prompt,
                    raw_session_history=raw_session_history,
                    pydantic_messages_count=len(pydantic_message_history),
                )
        except Exception as e:
            logger.warning(f"ask_agent '{agent_name}': failed to load session history: {e}")
            # Fall back to running without history

    # Run agent with timeout and context propagation
    logger.info(f"Invoking agent '{agent_name}' with prompt: {prompt[:100]}...")

    # Check if we have an event sink for streaming
    push_event = get_event_sink()
    use_streaming = push_event is not None

    streamed_content = ""  # Track if content was streamed

    try:
        # Set child context for nested tool calls
        with agent_context_scope(child_context):
            if use_streaming:
                # STREAMING MODE: Use iter() and proxy events to parent
                logger.debug(f"ask_agent '{agent_name}': using streaming mode with event proxying")

                async def run_with_streaming():
                    from pydantic_ai.messages import (
                        PartStartEvent, PartDeltaEvent, PartEndEvent,
                        FunctionToolResultEvent, FunctionToolCallEvent,
                    )
                    from pydantic_ai.agent import Agent

                    accumulated_content = []
                    child_tool_calls = []

                    # iter() returns an async context manager, not an awaitable
                    iter_kwargs = {"message_history": pydantic_message_history} if pydantic_message_history else {}
                    async with agent_runtime.iter(prompt, **iter_kwargs) as agent_run:
                        async for node in agent_run:
                            if Agent.is_model_request_node(node):
                                async with node.stream(agent_run.ctx) as request_stream:
                                    async for event in request_stream:
                                        # Proxy part starts (text content only - tool calls handled in is_call_tools_node)
                                        if isinstance(event, PartStartEvent):
                                            from pydantic_ai.messages import ToolCallPart, TextPart
                                            if isinstance(event.part, ToolCallPart):
                                                # Track tool call for later (args are incomplete at PartStartEvent)
                                                # Full args come via FunctionToolCallEvent in is_call_tools_node
                                                child_tool_calls.append({
                                                    "tool_name": event.part.tool_name,
                                                    "index": event.index,
                                                })
                                            elif isinstance(event.part, TextPart):
                                                # TextPart may have initial content
                                                if event.part.content:
                                                    accumulated_content.append(event.part.content)
                                                    await push_event.put({
                                                        "type": "child_content",
                                                        "agent_name": agent_name,
                                                        "content": event.part.content,
                                                    })
                                        # Proxy text content deltas to parent for real-time streaming
                                        elif isinstance(event, PartDeltaEvent):
                                            if hasattr(event, 'delta') and hasattr(event.delta, 'content_delta'):
                                                content = event.delta.content_delta
                                                if content:
                                                    accumulated_content.append(content)
                                                    # Push content chunk to parent for streaming
                                                    await push_event.put({
                                                        "type": "child_content",
                                                        "agent_name": agent_name,
                                                        "content": content,
                                                    })

                            elif Agent.is_call_tools_node(node):
                                async with node.stream(agent_run.ctx) as tools_stream:
                                    async for tool_event in tools_stream:
                                        # FunctionToolCallEvent fires when tool call is parsed
                                        # with complete arguments (before execution)
                                        if isinstance(tool_event, FunctionToolCallEvent):
                                            # Get full arguments from completed tool call
                                            tool_args = None
                                            if hasattr(tool_event, 'part') and hasattr(tool_event.part, 'args'):
                                                raw_args = tool_event.part.args
                                                if isinstance(raw_args, str):
                                                    try:
                                                        tool_args = json.loads(raw_args)
                                                    except json.JSONDecodeError:
                                                        tool_args = {"raw": raw_args}
                                                elif isinstance(raw_args, dict):
                                                    tool_args = raw_args
                                            # Push tool start with full arguments
                                            await push_event.put({
                                                "type": "child_tool_start",
                                                "agent_name": agent_name,
                                                "tool_name": tool_event.part.tool_name if hasattr(tool_event, 'part') else "unknown",
                                                "arguments": tool_args,
                                            })
                                        elif isinstance(tool_event, FunctionToolResultEvent):
                                            result_content = tool_event.result.content if hasattr(tool_event.result, 'content') else tool_event.result
                                            # Push tool result to parent
                                            await push_event.put({
                                                "type": "child_tool_result",
                                                "agent_name": agent_name,
                                                "result": result_content,
                                            })

                        # Get final result (inside context manager)
                        return agent_run.result, "".join(accumulated_content), child_tool_calls

                result, streamed_content, tool_calls = await asyncio.wait_for(
                    run_with_streaming(),
                    timeout=timeout_seconds
                )
            else:
                # NON-STREAMING MODE: Use run() for backwards compatibility
                if pydantic_message_history:
                    result = await asyncio.wait_for(
                        agent_runtime.run(prompt, message_history=pydantic_message_history),
                        timeout=timeout_seconds
                    )
                else:
                    result = await asyncio.wait_for(
                        agent_runtime.run(prompt),
                        timeout=timeout_seconds
                    )
    except asyncio.TimeoutError:
        return {
            "status": "error",
            "error": f"Agent '{agent_name}' timed out after {timeout_seconds}s",
            "agent_schema": agent_name,
        }

    # Serialize output
    from rem.agentic.serialization import serialize_agent_result, is_pydantic_model
    output = serialize_agent_result(result.output)

    logger.info(f"Agent '{agent_name}' completed successfully")

    # If child agent returned structured output (Pydantic model), emit as tool_call SSE event
    # This allows the frontend to render structured results (forms, cards, etc.)
    is_structured_output = is_pydantic_model(result.output)
    structured_tool_id = f"{agent_name}_structured_output"
    logger.debug(f"ask_agent '{agent_name}': is_structured_output={is_structured_output}, output_type={type(result.output).__name__}")

    if use_streaming and is_structured_output and push_event is not None:
        # Emit structured output as a tool_call event with the serialized result
        # Use agent_name as tool_name so it appears as the logical tool (e.g., "finalize_intake_agent")
        await push_event.put({
            "type": "tool_call",
            "tool_name": agent_name,  # Use agent name as tool name for clarity
            "tool_id": structured_tool_id,
            "status": "completed",
            "arguments": {"input_text": input_text},
            "result": output,  # Serialized Pydantic model as dict
        })
        logger.debug(f"ask_agent '{agent_name}': emitted structured output as tool_call SSE event")

    # Save structured output as a tool message in the database
    # This makes structured output agents look like tool calls in session history
    if is_structured_output and child_context and child_context.session_id and settings.postgres.enabled:
        try:
            from ...services.session import SessionMessageStore
            from ...utils.date_utils import utc_now, to_iso

            store = SessionMessageStore(user_id=child_context.user_id or "default")

            # Build tool message in the same format as regular tool calls
            tool_message = {
                "role": "tool",
                "content": json.dumps(output, default=str),  # Structured output as JSON
                "timestamp": to_iso(utc_now()),
                "tool_call_id": structured_tool_id,
                "tool_name": agent_name,  # Agent name as tool name
                "tool_arguments": {"input_text": input_text},
            }

            # Store as a single message (not using store_session_messages to avoid compression)
            await store.store_session_messages(
                session_id=child_context.session_id,
                messages=[tool_message],
                user_id=child_context.user_id,
                compress=False,  # Don't compress tool results
            )
            logger.debug(f"ask_agent '{agent_name}': saved structured output as tool message in session")
        except Exception as e:
            logger.warning(f"ask_agent '{agent_name}': failed to save structured output to database: {e}")

    response = {
        "status": "success",
        "output": output,
        "agent_schema": agent_name,
        "input_text": input_text,
        "is_structured_output": is_structured_output,  # Flag for caller to know result type
    }

    # Only include text_response if content was NOT streamed
    # When streaming, child_content events already delivered the content
    if not use_streaming or not streamed_content:
        response["text_response"] = str(result.output)

    return response


# =============================================================================
# Test/Debug Tools (for development only)
# =============================================================================

@mcp_tool_error_handler
async def test_error_handling(
    error_type: Literal["exception", "error_response", "timeout", "success"] = "success",
    delay_seconds: float = 0,
    error_message: str = "Test error occurred",
) -> dict[str, Any]:
    """
    Test tool for simulating different error scenarios.

    **FOR DEVELOPMENT/TESTING ONLY** - This tool helps verify that error
    handling works correctly through the streaming layer.

    Args:
        error_type: Type of error to simulate:
            - "success": Returns successful response (default)
            - "exception": Raises an exception (tests @mcp_tool_error_handler)
            - "error_response": Returns {"status": "error", ...} dict
            - "timeout": Delays for 60 seconds (simulates timeout)
        delay_seconds: Optional delay before responding (0-10 seconds)
        error_message: Custom error message for error scenarios

    Returns:
        Dict with test results or error information

    Examples:
        # Test successful response
        test_error_handling(error_type="success")

        # Test exception handling
        test_error_handling(error_type="exception", error_message="Database connection failed")

        # Test error response format
        test_error_handling(error_type="error_response", error_message="Resource not found")

        # Test with delay
        test_error_handling(error_type="success", delay_seconds=2)
    """
    import asyncio

    logger.info(f"test_error_handling called: type={error_type}, delay={delay_seconds}")

    # Apply delay (capped at 10 seconds for safety)
    if delay_seconds > 0:
        await asyncio.sleep(min(delay_seconds, 10))

    if error_type == "exception":
        # This tests the @mcp_tool_error_handler decorator
        raise RuntimeError(f"TEST EXCEPTION: {error_message}")

    elif error_type == "error_response":
        # This tests how the streaming layer handles error status responses
        return {
            "status": "error",
            "error": error_message,
            "error_code": "TEST_ERROR",
            "recoverable": True,
        }

    elif error_type == "timeout":
        # Simulate a very long operation (for testing client-side timeouts)
        await asyncio.sleep(60)
        return {"status": "success", "message": "Timeout test completed (should not reach here)"}

    else:  # success
        return {
            "status": "success",
            "message": "Test completed successfully",
            "test_data": {
                "error_type": error_type,
                "delay_applied": delay_seconds,
                "timestamp": str(asyncio.get_event_loop().time()),
            },
        }


# =============================================================================
# Vision Tools
# =============================================================================


@mcp_tool_error_handler
async def analyze_pages(
    file_uri: str,
    start_page: int = 1,
    end_page: int | None = None,
    prompt: str = "Extract all text and tables from these pages",
    provider: str = "anthropic",
    model: str | None = None,
    page_batch_size: int = 5,
) -> dict[str, Any]:
    """
    Analyze PDF or image pages using vision models via Pydantic AI.

    Batches pages together in single API calls for efficiency.
    Uses Pydantic AI for full OpenTelemetry tracing and cost tracking.

    Args:
        file_uri: Path to PDF or image file (local path, s3://, or https://)
        start_page: First page to analyze (1-indexed)
        end_page: Last page to analyze (inclusive, defaults to last page)
        prompt: Instruction for what to extract/analyze
        provider: Vision provider (anthropic, openai, gemini)
        model: Optional model override (e.g., "claude-sonnet-4.5", "gpt-4.1", "gpt-4o")
        page_batch_size: How many pages to send to the model at once (default: 5)

    Returns:
        Dict with vision analysis results including:
        - pages: List of page numbers processed
        - page_count: Number of pages processed
        - result: Vision analysis output text
        - provider: Provider used
        - model: Model identifier used
        - usage: Token usage statistics

    Examples:
        # Extract text from a PDF
        analyze_pages(
            file_uri="s3://bucket/document.pdf",
            prompt="Extract all text and tables"
        )

        # Analyze specific pages with GPT-4
        analyze_pages(
            file_uri="/path/to/invoice.pdf",
            start_page=1,
            end_page=3,
            prompt="Extract invoice line items as JSON",
            provider="openai",
            model="gpt-4o"
        )

        # Process an image
        analyze_pages(
            file_uri="/path/to/scan.png",
            prompt="Describe what you see in this image"
        )
    """
    from pathlib import Path
    from ...services.vision import analyze_images_async, VisionProvider, MIME_TYPES

    file_path = Path(file_uri)
    if not file_path.exists():
        return {"status": "error", "error": f"File not found: {file_uri}"}

    suffix = file_path.suffix.lower()

    # For single image files, just process directly
    if suffix in [".png", ".jpg", ".jpeg", ".webp", ".gif"]:
        image_bytes = file_path.read_bytes()
        media_type = MIME_TYPES.get(suffix, "image/png")
        try:
            provider_enum = VisionProvider(provider.lower())
            result = await analyze_images_async(
                images=[(image_bytes, media_type)],
                prompt=prompt,
                provider=provider_enum,
                model=model,
            )
            return {
                "status": "success",
                "pages": [1],
                "page_count": 1,
                "prompt": prompt,
                "result": result.description,
                "provider": provider,
                "model": result.model,
                "usage": result.usage,
            }
        except Exception as e:
            logger.error(f"Vision analysis failed: {e}")
            return {"status": "error", "error": f"Vision analysis failed: {e}"}

    # For PDFs, render pages to images using PyMuPDF
    if suffix != ".pdf":
        return {"status": "error", "error": f"Unsupported file type: {suffix}"}

    try:
        import fitz  # PyMuPDF

        doc = fitz.open(file_path)
        total_pages = len(doc)

        # Determine page range
        if end_page is None:
            end_page = total_pages

        # Validate page range (1-indexed input, 0-indexed internal)
        start_page = max(1, min(start_page, total_pages))
        end_page = max(start_page, min(end_page, total_pages))

        # Render requested pages to images
        page_images = []
        for page_num in range(start_page - 1, end_page):  # Convert to 0-indexed
            page = doc[page_num]
            # Render at 150 DPI for good quality while keeping size reasonable
            mat = fitz.Matrix(150 / 72, 150 / 72)
            pix = page.get_pixmap(matrix=mat)
            image_bytes = pix.tobytes("png")
            page_images.append((page_num + 1, image_bytes))  # Store 1-indexed page number

        doc.close()

        if not page_images:
            return {
                "status": "error",
                "error": f"No pages found in range {start_page}-{end_page}",
                "total_pages": total_pages
            }

        logger.info(f"analyze_pages: Processing {len(page_images)} pages in batches of {page_batch_size}")

        # Process pages in batches
        all_results = []
        all_usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "requests": 0}
        provider_enum = VisionProvider(provider.lower())
        model_used = None
        pages_processed = []

        for batch_start in range(0, len(page_images), page_batch_size):
            batch = page_images[batch_start:batch_start + page_batch_size]
            batch_pages = [page_num for page_num, _ in batch]

            # Prepare images for batch
            images_data = [(img_bytes, "image/png") for _, img_bytes in batch]

            # Add page context to prompt for multi-page batches
            batch_prompt = prompt
            if len(batch) > 1:
                batch_prompt = f"Pages {batch_pages[0]}-{batch_pages[-1]}: {prompt}"

            try:
                batch_result = await analyze_images_async(
                    images=images_data,
                    prompt=batch_prompt,
                    provider=provider_enum,
                    model=model,
                )

                all_results.append(batch_result.description)
                model_used = batch_result.model
                pages_processed.extend(batch_pages)

                # Aggregate usage
                if batch_result.usage:
                    all_usage["input_tokens"] += batch_result.usage.get("input_tokens", 0)
                    all_usage["output_tokens"] += batch_result.usage.get("output_tokens", 0)
                    all_usage["total_tokens"] += batch_result.usage.get("total_tokens", 0)
                    all_usage["requests"] += batch_result.usage.get("requests", 0)

                logger.info(f"analyze_pages: Batch complete, pages {batch_pages}")

            except Exception as e:
                logger.error(f"Batch analysis failed for pages {batch_pages}: {e}")
                all_results.append(f"[Error processing pages {batch_pages}: {e}]")

        # Combine results
        combined_result = "\n\n---\n\n".join(all_results) if len(all_results) > 1 else all_results[0] if all_results else ""

        return {
            "status": "success",
            "pages": pages_processed,
            "page_count": len(pages_processed),
            "batch_size": page_batch_size,
            "batches_processed": (len(page_images) + page_batch_size - 1) // page_batch_size,
            "prompt": prompt,
            "result": combined_result,
            "provider": provider,
            "model": model_used,
            "usage": all_usage,
        }

    except ImportError:
        return {"status": "error", "error": "PyMuPDF (fitz) not available. Install with: pip install pymupdf"}
    except Exception as e:
        logger.error(f"Vision analysis failed: {e}")
        return {"status": "error", "error": f"Vision analysis failed: {e}"}
