"""
Recce MCP (Model Context Protocol) Server

This module implements a stdio-based MCP server that provides tools for
interacting with Recce's data validation capabilities.
"""

import asyncio
import json
import logging
import os
import textwrap
import time
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import TextContent, Tool

from recce.core import RecceContext, load_context
from recce.exceptions import RecceException
from recce.server import RecceServerMode
from recce.tasks.dataframe import DataFrame
from recce.tasks.profile import ProfileDiffTask
from recce.tasks.query import QueryDiffTask, QueryTask
from recce.tasks.rowcount import RowCountDiffTask

logger = logging.getLogger(__name__)


def _truncate_strings(obj: Any, max_length: int = 200) -> Any:
    """Recursively truncate strings longer than max_length in nested dicts and lists"""
    if isinstance(obj, dict):
        return {k: _truncate_strings(v, max_length) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [_truncate_strings(item, max_length) for item in obj]
    elif isinstance(obj, str) and len(obj) > max_length:
        return obj[:max_length] + "..."
    return obj


class MCPLogger:
    """JSON logger for MCP server request/response logging"""

    def __init__(self, debug: bool = False, log_file: str = "logs/recce-mcp.json"):
        self.debug = debug
        self.log_file = log_file

        if self.debug:
            # Create logs directory if it doesn't exist
            log_dir = os.path.dirname(log_file)
            if log_dir:
                os.makedirs(log_dir, exist_ok=True)

            # Overwrite log file on initialization
            try:
                with open(log_file, "w") as f:
                    f.write("")  # Clear existing content
            except Exception as e:
                logger.warning(f"Failed to initialize log file {log_file}: {e}")

    def _write_log(self, log_entry: Dict[str, Any]) -> None:
        """Write a log entry to the JSON file"""
        if not self.debug:
            return

        try:
            with open(self.log_file, "a") as f:
                f.write(json.dumps(log_entry) + "\n")
        except Exception as e:
            logger.warning(f"Failed to write to log file {self.log_file}: {e}")

    def log_list_tools(self, tools: List[Tool]) -> None:
        """Log a list_tools call"""
        tool_names = [tool.name for tool in tools]
        log_entry = {
            "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
            "type": "list_tools",
            "tools": tool_names,
        }
        self._write_log(log_entry)

    def log_tool_call(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        response: Dict[str, Any],
        duration_ms: float,
        error: Optional[str] = None,
    ) -> None:
        """Log a tool call with request and response"""
        log_entry = {
            "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
            "type": "call_tool",
            "tool": tool_name,
            "request": arguments,
            "duration_ms": round(duration_ms, 2),
        }

        if error:
            log_entry["error"] = error
        else:
            log_entry["response"] = _truncate_strings(response)

        self._write_log(log_entry)


class RecceMCPServer:
    """MCP Server for Recce data validation tools"""

    def __init__(
        self,
        context: RecceContext,
        mode: Optional[RecceServerMode] = None,
        debug: bool = False,
        log_file: str = "logs/recce-mcp.json",
        state_loader=None,
    ):
        self.context = context
        self.state_loader = state_loader
        self.mode = mode or RecceServerMode.server
        self.server = Server("recce")
        self.mcp_logger = MCPLogger(debug=debug, log_file=log_file)
        self._setup_handlers()

    def _setup_handlers(self):
        """Register all tool handlers"""

        @self.server.list_tools()
        async def list_tools() -> List[Tool]:
            """List all available tools based on server mode"""
            logger.info(f"[MCP] list_tools called (mode: {self.mode.value if self.mode else 'server'})")
            tools = []

            # Always available in all modes
            tools.append(
                Tool(
                    name="lineage_diff",
                    description=textwrap.dedent(
                        """
                        Get the lineage diff between production(base) and session(current) for changed models.
                        Returns nodes and edges (node dependencies) in compact dataframe format.

                        Nodes dataframe includes: idx, id, name, resource_type, materialized, change_status, impacted.
                        Edges dataframe includes: from (parent node idx), to (child node idx).

                        Rendering guidance for Mermaid diagram:
                        Use graph LR and apply these styles based on change_status and impacted:
                        - change_status="added": fill:#d4edda, stroke:#28a745, color:#000000
                        - change_status="removed": fill:#f8d7da, stroke:#dc3545, color:#000000
                        - change_status="modified" AND impacted=true: fill:#fff3cd, stroke:#ffc107, color:#000000
                        - change_status=null AND impacted=true: fill:#ffffff, stroke:#ffc107, color:#000000
                        - change_status=null AND impacted=false: fill:#ffffff, stroke:#d3d3d3, color:#999999
                    """
                    ).strip(),
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "select": {
                                "type": "string",
                                "description": "dbt selector syntax to filter models (optional)",
                            },
                            "exclude": {
                                "type": "string",
                                "description": "dbt selector syntax to exclude models (optional)",
                            },
                            "packages": {
                                "type": "array",
                                "items": {"type": "string"},
                                "description": "List of packages to filter (optional)",
                            },
                            "view_mode": {
                                "type": "string",
                                "enum": ["changed_models", "all"],
                                "default": "changed_models",
                                "description": "View mode: 'changed_models' for only changed models (default), 'all' for all models",
                            },
                        },
                    },
                )
            )
            tools.append(
                Tool(
                    name="schema_diff",
                    description="Get the schema diff (column changes) between base and current environments. "
                    "Shows added, removed, and type-changed columns in compact dataframe format.",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "select": {
                                "type": "string",
                                "description": "dbt selector syntax to filter models (optional)",
                            },
                            "exclude": {
                                "type": "string",
                                "description": "dbt selector syntax to exclude models (optional)",
                            },
                            "packages": {
                                "type": "array",
                                "items": {"type": "string"},
                                "description": "List of packages to filter (optional)",
                            },
                        },
                    },
                )
            )

            # Diff tools only available in server mode, not in preview or read-only mode
            if self.mode == RecceServerMode.server:
                tools.extend(
                    [
                        Tool(
                            name="row_count_diff",
                            description="Compare row counts between base and current environments for specified models.",
                            inputSchema={
                                "type": "object",
                                "properties": {
                                    "node_names": {
                                        "type": "array",
                                        "items": {"type": "string"},
                                        "description": "List of model names to check row counts (optional)",
                                    },
                                    "node_ids": {
                                        "type": "array",
                                        "items": {"type": "string"},
                                        "description": "List of node IDs to check row counts (optional)",
                                    },
                                    "select": {
                                        "type": "string",
                                        "description": "dbt selector syntax to filter models (optional)",
                                    },
                                    "exclude": {
                                        "type": "string",
                                        "description": "dbt selector syntax to exclude models (optional)",
                                    },
                                },
                            },
                        ),
                        Tool(
                            name="query",
                            description="Execute a SQL query on the current environment. "
                            "Supports Jinja templates with dbt macros like {{ ref('model_name') }}.",
                            inputSchema={
                                "type": "object",
                                "properties": {
                                    "sql_template": {
                                        "type": "string",
                                        "description": "SQL query template with optional Jinja syntax",
                                    },
                                    "base": {
                                        "type": "boolean",
                                        "description": "Whether to run on base environment (default: false)",
                                        "default": False,
                                    },
                                },
                                "required": ["sql_template"],
                            },
                        ),
                        Tool(
                            name="query_diff",
                            description="Execute SQL queries on both base and current environments and compare results. "
                            "Supports primary keys for row-level comparison.",
                            inputSchema={
                                "type": "object",
                                "properties": {
                                    "sql_template": {
                                        "type": "string",
                                        "description": "SQL query template for current environment",
                                    },
                                    "base_sql_template": {
                                        "type": "string",
                                        "description": "SQL query template for base environment (optional, defaults to sql_template)",
                                    },
                                    "primary_keys": {
                                        "type": "array",
                                        "items": {"type": "string"},
                                        "description": "List of primary key columns for row comparison (optional)",
                                    },
                                },
                                "required": ["sql_template"],
                            },
                        ),
                        Tool(
                            name="profile_diff",
                            description="Generate and compare statistical profiles (min, max, avg, distinct count, etc.) "
                            "for columns in a model between base and current environments.",
                            inputSchema={
                                "type": "object",
                                "properties": {
                                    "model": {
                                        "type": "string",
                                        "description": "Model name to profile",
                                    },
                                    "columns": {
                                        "type": "array",
                                        "items": {"type": "string"},
                                        "description": "List of column names to profile (optional, profiles all columns if not specified)",
                                    },
                                },
                                "required": ["model"],
                            },
                        ),
                        Tool(
                            name="list_checks",
                            description="List all checks in the current session. Returns check metadata including check IDs, names, types, parameters, and approval status.",
                            inputSchema={
                                "type": "object",
                                "properties": {},
                            },
                        ),
                        Tool(
                            name="run_check",
                            description="Run a single check by ID and wait for completion. Returns execution status, results, and approval status.",
                            inputSchema={
                                "type": "object",
                                "properties": {
                                    "check_id": {
                                        "type": "string",
                                        "description": "The ID of the check to run",
                                    },
                                },
                                "required": ["check_id"],
                            },
                        ),
                    ]
                )

            self.mcp_logger.log_list_tools(tools)

            # Log available tools to console
            tool_names = [tool.name for tool in tools]
            logger.info(f"[MCP] Returning {len(tools)} tools: {', '.join(tool_names)}")

            return tools

        @self.server.call_tool()
        async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
            """Handle tool calls"""
            start_time = time.perf_counter()

            # Log incoming request
            logger.info(f"[MCP] Tool call received: {name}")
            logger.info(f"[MCP] Arguments: {json.dumps(arguments, indent=2)}")

            try:
                # Check if tool is blocked in non-server mode
                blocked_tools_in_non_server = {
                    "row_count_diff",
                    "query",
                    "query_diff",
                    "profile_diff",
                    "list_checks",
                    "run_check",
                }
                if self.mode != RecceServerMode.server and name in blocked_tools_in_non_server:
                    raise ValueError(
                        f"Tool '{name}' is not available in {self.mode.value} mode. "
                        "Only 'lineage_diff' and 'schema_diff' are available in this mode."
                    )

                if name == "lineage_diff":
                    result = await self._tool_lineage_diff(arguments)
                elif name == "schema_diff":
                    result = await self._tool_schema_diff(arguments)
                elif name == "row_count_diff":
                    result = await self._tool_row_count_diff(arguments)
                elif name == "query":
                    result = await self._tool_query(arguments)
                elif name == "query_diff":
                    result = await self._tool_query_diff(arguments)
                elif name == "profile_diff":
                    result = await self._tool_profile_diff(arguments)
                elif name == "list_checks":
                    result = await self._tool_list_checks(arguments)
                elif name == "run_check":
                    result = await self._tool_run_check(arguments)
                else:
                    raise ValueError(f"Unknown tool: {name}")

                duration_ms = (time.perf_counter() - start_time) * 1000
                self.mcp_logger.log_tool_call(name, arguments, result, duration_ms)

                # Log outgoing response
                response_json = json.dumps(result, indent=2)
                logger.info(f"[MCP] Tool response for {name} ({duration_ms:.2f}ms):")
                # Truncate large responses for console readability
                if len(response_json) > 1000:
                    logger.debug(f"[MCP] {response_json[:1000]}... (truncated, {len(response_json)} chars total)")
                else:
                    logger.debug(f"[MCP] {response_json}")

                return [TextContent(type="text", text=response_json)]
            except Exception as e:
                duration_ms = (time.perf_counter() - start_time) * 1000
                self.mcp_logger.log_tool_call(name, arguments, {}, duration_ms, error=str(e))
                logger.error(f"[MCP] Error executing tool {name} ({duration_ms:.2f}ms): {str(e)}")
                logger.exception("[MCP] Full traceback:")
                error_response = json.dumps({"error": str(e)}, indent=2)
                return [TextContent(type="text", text=error_response)]

    async def _tool_lineage_diff(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Get lineage diff between base and current"""
        try:
            # Extract filter arguments
            select = arguments.get("select")
            exclude = arguments.get("exclude")
            packages = arguments.get("packages")
            view_mode = arguments.get("view_mode", "changed_models")

            # Get lineage diff from adapter (returns a Pydantic LineageDiff model)
            lineage_diff = self.context.get_lineage_diff().model_dump(mode="json")

            # Apply node selection filtering if arguments provided
            selected_node_ids = self.context.adapter.select_nodes(
                select=select,
                exclude=exclude,
                packages=packages,
                view_mode=view_mode,
            )
            impacted_node_ids = self.context.adapter.select_nodes(
                select="state:modified+",
            )

            # Get diff information for change_status
            diff_info = lineage_diff.get("diff", {})

            # Extract parent_map and simplified nodes from both base and current
            parent_map = {}
            nodes = {}

            # Merge parent_map and nodes: base first, then current overrides
            for env_key in ["base", "current"]:
                if env_key not in lineage_diff:
                    continue

                env_data = lineage_diff[env_key]

                # Merge parent_map (filtering by selected nodes)
                if "parent_map" in env_data:
                    for node_id, parents in env_data["parent_map"].items():
                        if node_id in selected_node_ids:
                            parent_map[node_id] = parents

                # Merge nodes (filtering by selected nodes)
                if "nodes" in env_data:
                    for node_id, node_info in env_data["nodes"].items():
                        if node_id in selected_node_ids:
                            nodes[node_id] = {
                                "name": node_info.get("name"),
                                "resource_type": node_info.get("resource_type"),
                            }

                            materialized = node_info.get("config", {}).get("materialized")
                            if materialized is not None:
                                nodes[node_id]["materialized"] = materialized

            # Create id to idx mapping
            id_to_idx = {node_id: idx for idx, node_id in enumerate(nodes.keys())}

            # Prepare node data for DataFrame
            nodes_data = [
                (
                    id_to_idx[node_id],
                    node_id,
                    node_info.get("name"),
                    node_info.get("resource_type"),
                    node_info.get("materialized"),
                    diff_info.get(node_id, {}).get("change_status"),
                    node_id in impacted_node_ids,
                )
                for node_id, node_info in nodes.items()
            ]

            # Create nodes DataFrame using from_data with simple dict format
            nodes_df = DataFrame.from_data(
                columns={
                    "idx": "integer",
                    "id": "text",
                    "name": "text",
                    "resource_type": "text",
                    "materialized": "text",
                    "change_status": "text",
                    "impacted": "boolean",
                },
                data=nodes_data,
            )

            # Build edges from parent_map
            edges_data = []
            for node_id, parents in parent_map.items():
                if node_id in id_to_idx:
                    for parent_id in parents:
                        if parent_id in id_to_idx:
                            edges_data.append((id_to_idx[parent_id], id_to_idx[node_id]))

            # Create edges DataFrame
            edges_df = DataFrame.from_data(
                columns={
                    "from": "integer",
                    "to": "integer",
                },
                data=edges_data,
            )

            # Build simplified result
            result = {"nodes": nodes_df.model_dump(mode="json"), "edges": edges_df.model_dump(mode="json")}

            return result

        except Exception:
            logger.exception("Error getting lineage diff")
            raise

    async def _tool_schema_diff(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Get schema diff (column changes) between base and current"""
        try:
            # Extract filter arguments
            select = arguments.get("select")
            exclude = arguments.get("exclude")
            packages = arguments.get("packages")

            # Get lineage diff from adapter
            lineage_diff = self.context.get_lineage_diff().model_dump(mode="json")

            # Get all nodes from current environment
            current_nodes = {}
            if "current" in lineage_diff and "nodes" in lineage_diff["current"]:
                current_nodes = lineage_diff["current"]["nodes"]

            # Filter to only nodes that exist in both base and current (exclude added nodes)
            base_nodes = lineage_diff.get("base", {}).get("nodes", {})
            nodes_to_compare = set(current_nodes.keys()) & set(base_nodes.keys())

            # Apply filtering if arguments provided
            if select or exclude or packages:
                selected_node_ids = self.context.adapter.select_nodes(
                    select=select,
                    exclude=exclude,
                    packages=packages,
                )
                nodes_to_compare = nodes_to_compare & selected_node_ids

            # Build schema changes
            schema_changes = []

            for node_id in nodes_to_compare:
                base_node = base_nodes.get(node_id, {})
                current_node = current_nodes.get(node_id, {})

                base_columns = base_node.get("columns", {})
                current_columns = current_node.get("columns", {})

                # Get column names in base and current
                base_col_names = set(base_columns.keys())
                current_col_names = set(current_columns.keys())

                # Find added columns (in current but not in base)
                for col_name in current_col_names - base_col_names:
                    schema_changes.append((node_id, col_name, "added"))

                # Find removed columns (in base but not in current)
                for col_name in base_col_names - current_col_names:
                    schema_changes.append((node_id, col_name, "removed"))

                # Find modified columns (in both but with different types)
                for col_name in base_col_names & current_col_names:
                    base_col_type = base_columns[col_name].get("type")
                    current_col_type = current_columns[col_name].get("type")
                    if base_col_type != current_col_type:
                        schema_changes.append((node_id, col_name, "modified"))

            # Check if there are more than 100 rows
            limit = 100
            has_more = len(schema_changes) > limit
            limited_schema_changes = schema_changes[:limit]

            # Convert schema changes to dataframe format using DataFrame.from_data()
            diff_df = DataFrame.from_data(
                columns={
                    "node_id": "text",
                    "column": "text",
                    "change_status": "text",
                },
                data=limited_schema_changes,
                limit=limit,
                more=has_more,
            )
            return diff_df.model_dump(mode="json")

        except Exception:
            logger.exception("Error getting schema diff")
            raise

    async def _tool_row_count_diff(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Execute row count diff task"""
        try:
            task = RowCountDiffTask(params=arguments)

            # Execute task synchronously (it's already sync)
            result = await asyncio.get_event_loop().run_in_executor(None, task.execute)

            return result
        except Exception:
            logger.exception("Error executing row count diff")
            raise

    async def _tool_query(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a query"""
        try:
            sql_template = arguments.get("sql_template")
            is_base = arguments.get("base", False)

            params = {"sql_template": sql_template}
            task = QueryTask(params=params)
            task.is_base = is_base

            # Execute task
            result = await asyncio.get_event_loop().run_in_executor(None, task.execute)

            # Convert to dict if it's a model
            if hasattr(result, "model_dump"):
                return result.model_dump(mode="json")
            return result
        except Exception:
            logger.exception("Error executing query")
            raise

    async def _tool_query_diff(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Execute query diff task"""
        try:
            task = QueryDiffTask(params=arguments)

            # Execute task
            result = await asyncio.get_event_loop().run_in_executor(None, task.execute)

            # Convert to dict if it's a model
            if hasattr(result, "model_dump"):
                return result.model_dump(mode="json")
            return result
        except Exception:
            logger.exception("Error executing query diff")
            raise

    async def _tool_profile_diff(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Execute profile diff task"""
        try:
            task = ProfileDiffTask(params=arguments)

            # Execute task
            result = await asyncio.get_event_loop().run_in_executor(None, task.execute)

            # Convert to dict if it's a model
            if hasattr(result, "model_dump"):
                return result.model_dump(mode="json")
            return result
        except Exception:
            logger.exception("Error executing profile diff")
            raise

    async def _tool_list_checks(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """List all checks in the current session"""
        try:
            from recce.models import CheckDAO

            # Get all checks
            checks_dao = CheckDAO()
            checks = checks_dao.list()

            # Build checks list with relevant metadata
            checks_list = []
            for check in checks:
                checks_list.append(
                    {
                        "check_id": str(check.check_id),
                        "name": check.name,
                        "type": check.type.value,
                        "description": check.description or "",
                        "params": check.params or {},
                        "is_checked": check.is_checked,
                        "is_preset": check.is_preset,
                    }
                )

            # Get statistics
            stats = checks_dao.status()

            result = {
                "checks": checks_list,
                "total": stats["total"],
                "approved": stats["approved"],
            }

            return result
        except Exception:
            logger.exception("Error listing checks")
            raise

    async def _tool_run_check(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Run a single check by ID"""
        try:
            from recce.apis.run_func import submit_run
            from recce.models import CheckDAO
            from recce.models.types import RunType

            check_id = arguments.get("check_id")
            if not check_id:
                raise ValueError("check_id is required")

            # Get the check from database
            check = CheckDAO().find_check_by_id(check_id)
            if not check:
                raise ValueError(f"Check with ID {check_id} not found")

            # For lineage_diff and schema_diff, call the corresponding diff tools
            if check.type == RunType.LINEAGE_DIFF:
                return await self._tool_lineage_diff(check.params or {})

            if check.type == RunType.SCHEMA_DIFF:
                return await self._tool_schema_diff(check.params or {})

            try:
                # Submit and execute the run for other check types
                run, future = submit_run(check.type, params=check.params or {}, check_id=check_id)
                run.result = await future

                # Return run object in JSON format (same as run_check_handler in check_api.py)
                return run.model_dump(mode="json")
            except RecceException as e:
                raise ValueError(str(e))

        except Exception:
            logger.exception("Error running check")
            raise

    async def run(self):
        """Run the MCP server in stdio mode"""
        try:
            async with stdio_server() as (read_stream, write_stream):
                await self.server.run(read_stream, write_stream, self.server.create_initialization_options())
        finally:
            # Export state on shutdown if state_loader is available
            if self.state_loader and self.context:
                try:
                    from rich.console import Console

                    console = Console()

                    # Export the state
                    msg = self.state_loader.export(self.context.export_state())
                    if msg is not None:
                        console.print(f"[yellow]On shutdown:[/yellow] {msg}")
                    else:
                        if hasattr(self.state_loader, "state_file") and self.state_loader.state_file:
                            console.print(
                                f"[yellow]On shutdown:[/yellow] State exported to '{self.state_loader.state_file}'"
                            )
                        else:
                            console.print("[yellow]On shutdown:[/yellow] State exported successfully")
                except Exception as e:
                    logger.exception(f"Failed to export state on shutdown: {e}")

    async def run_sse(self, host: str = "localhost", port: int = 8000):
        """Run the MCP server in HTTP mode using Server-Sent Events (SSE)

        Args:
            host: Host to bind to (default: localhost)
            port: Port to bind to (default: 8000)
        """
        from contextlib import asynccontextmanager

        import uvicorn
        from mcp.server.sse import SseServerTransport
        from starlette.applications import Starlette
        from starlette.requests import Request
        from starlette.responses import Response
        from starlette.routing import Mount, Route

        # Create SSE transport - endpoint where clients POST messages
        sse = SseServerTransport("/")

        async def handle_sse_request(request: Request):
            """Handle SSE connection (GET /sse) following official MCP example"""
            client_info = f"{request.client.host}:{request.client.port}" if request.client else "unknown"
            logger.info(f"[MCP HTTP] SSE connection established from {client_info}")
            try:
                async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
                    await self.server.run(streams[0], streams[1], self.server.create_initialization_options())
            finally:
                logger.info(f"[MCP HTTP] SSE connection closed from {client_info}")
            return Response()  # Required to avoid NoneType error

        async def handle_post_message(scope, receive, send):
            """Handle POST messages (POST /) for MCP protocol"""
            # Log POST message (session_id will be in query params)
            query_string = scope.get("query_string", b"").decode("utf-8")
            logger.debug(f"[MCP HTTP] POST message received with query: {query_string}")
            await sse.handle_post_message(scope, receive, send)

        async def handle_health_check(request: Request):
            """Handle health check endpoint (GET /health)"""
            return Response(content='{"status":"ok"}', media_type="application/json")

        @asynccontextmanager
        async def lifespan(app):
            """Handle startup and shutdown events"""
            # Startup
            yield
            # Shutdown - this runs when server exits (SIGINT, SIGTERM, etc.)
            if self.state_loader and self.context:
                try:
                    logger.info("Exporting state on shutdown...")
                    msg = self.state_loader.export(self.context.export_state())
                    if msg:
                        logger.info(f"State export: {msg}")
                except Exception as e:
                    logger.exception(f"Failed to export state on shutdown: {e}")

        # Create Starlette app with lifespan
        app = Starlette(
            debug=self.mcp_logger.debug,
            routes=[
                Route("/health", endpoint=handle_health_check, methods=["GET"]),
                Route("/sse", endpoint=handle_sse_request, methods=["GET"]),
                Mount("/", app=handle_post_message),
            ],
            lifespan=lifespan,
        )

        # Run with uvicorn
        logger.info(f"Starting Recce MCP Server in HTTP mode on {host}:{port}")
        logger.info(f"Connection URL: http://{host}:{port}/sse")
        config = uvicorn.Config(app, host=host, port=port, log_level="info")
        server = uvicorn.Server(config)
        await server.serve()


async def run_mcp_server(
    sse: bool = False,
    host: str = "localhost",
    port: int = 8000,
    **kwargs,
):
    """
    Entry point for running the MCP server

    Args:
        sse: Whether to run in HTTP/SSE mode (default: False for stdio mode)
        host: Host to bind to in SSE mode (default: localhost)
        port: Port to bind to in SSE mode (default: 8000)
        **kwargs: Arguments for loading RecceContext (dbt options, etc.)
               Optionally includes 'state_loader' for persisting state on shutdown
               Optionally includes 'mode' for server mode (server, preview, read-only)
               Optionally includes 'debug' flag for enabling MCP logging
    """
    state_loader = kwargs.get("state_loader", None)
    # Setup logging
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")

    # Load Recce context
    context = load_context(**kwargs)

    # Extract mode from kwargs (defaults to server mode)
    mode_str = kwargs.get("mode")
    mode = None
    if mode_str:
        # Convert string mode to RecceServerMode enum
        try:
            mode = RecceServerMode(mode_str)
        except ValueError:
            logger.warning(f"Invalid mode '{mode_str}', using default server mode")

    # Extract debug flag from kwargs
    debug = kwargs.get("debug", False)

    # Create MCP server with state_loader for graceful shutdown
    server = RecceMCPServer(context, mode=mode, debug=debug, state_loader=state_loader)

    # Run in either stdio or SSE mode
    if sse:
        # SSE mode: lifespan handler in Starlette manages shutdown and state export
        await server.run_sse(host=host, port=port)
    else:
        # Stdio mode: run() method handles shutdown and state export via try-finally
        await server.run()
