"""
Sound Recording MCP Server

Exposes tools for querying sound recording data, fingerprints, collections,
and integrating with AWS services (Lambda, Step Functions, S3).
"""

import argparse
import asyncio
import json
import logging
import os
import sys
from pathlib import Path
from typing import Any, Dict

from dotenv import load_dotenv
from mcp.server import Server
from mcp.server.sse import SseServerTransport
from mcp.server.stdio import stdio_server
from mcp.types import TextContent, Tool


# Loading env variables
REPO_ROOT = Path(__file__).resolve().parent.parent
SRC_DIR = REPO_ROOT / "src"
if str(SRC_DIR) not in sys.path:
    sys.path.insert(0, str(SRC_DIR))
load_dotenv(REPO_ROOT / ".env")


from tools import mysql_tools, aws_tools, sound_recording_tools, neo4j_tools, snowflake_tools

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    stream=sys.stderr
)
logger = logging.getLogger(__name__)


def _json_default(value: Any) -> str:
    """Serialize non-standard values (e.g. neo4j.time.DateTime) as strings."""
    if hasattr(value, "iso_format"):
        return value.iso_format()
    return str(value)


# Initialize MCP server
server = Server("sound-recordings-mcp")


def _normalize_path(path: str) -> str:
    """Normalize a route path so it always starts with `/` and never ends with `/`."""
    if not path.startswith("/"):
        path = "/" + path
    return path.rstrip("/") or "/"


async def _send_json_response(send, status: int, payload: dict[str, Any]) -> None:
    body = json.dumps(payload, indent=2).encode("utf-8")
    headers = [
        (b"content-type", b"application/json; charset=utf-8"),
        (b"content-length", str(len(body)).encode("ascii")),
    ]
    await send({"type": "http.response.start", "status": status, "headers": headers})
    await send({"type": "http.response.body", "body": body})


def _build_sse_app(
        *,
        sse_path: str,
        message_path: str,
) -> Any:
    """Build a minimal ASGI app that serves MCP over SSE."""
    sse_path = _normalize_path(sse_path)
    message_path = _normalize_path(message_path)
    transport = SseServerTransport(message_path)
    initialization_options = server.create_initialization_options()

    async def app(scope, receive, send):
        if scope["type"] != "http":
            return

        path = _normalize_path(scope.get("path", "/"))
        method = scope.get("method", "GET").upper()

        if path == "/health" and method == "GET":
            await _send_json_response(
                send,
                200,
                {
                    "status": "ok",
                    "transport": "sse",
                    "sse_path": sse_path,
                    "message_path": message_path,
                },
            )
            return

        if path == sse_path and method == "GET":
            async with transport.connect_sse(scope, receive, send) as (
                    read_stream,
                    write_stream,
            ):
                await server.run(read_stream, write_stream, initialization_options)
            return

        if path == message_path and method == "POST":
            await transport.handle_post_message(scope, receive, send)
            return

        await _send_json_response(
            send,
            404,
            {
                "error": "Not found",
                "path": path,
                "available": ["GET /health", f"GET {sse_path}", f"POST {message_path}"],
            },
        )

    return app


async def _run_stdio_server() -> None:
    """Run the MCP server over stdio."""
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Sound Recording MCP Server")
    parser.add_argument(
        "--transport",
        choices=("stdio", "sse"),
        default=os.getenv("MCP_TRANSPORT", "stdio"),
        help="Transport to use (default: stdio)",
    )
    parser.add_argument(
        "--host",
        default=os.getenv("MCP_HOST", "127.0.0.1"),
        help="Host to bind when using HTTP transports",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=int(os.getenv("MCP_PORT", "55392")),
        help="Port to bind when using HTTP transports",
    )
    parser.add_argument(
        "--sse-path",
        default=os.getenv("MCP_SSE_PATH", "/sse"),
        help="GET endpoint for SSE connections",
    )
    parser.add_argument(
        "--message-path",
        default=os.getenv("MCP_MESSAGE_PATH", "/messages"),
        help="POST endpoint for SSE client messages",
    )
    return parser.parse_args()


@server.list_tools()
async def list_tools() -> list[Tool]:
    """Return list of available tools."""
    return [
        # Sound Recording Tools
        Tool(
            name="get_delivery_history",
            description="Get delivery history for sound recordings",
            inputSchema={
                "type": "object",
                "properties": {
                    "env": {
                        "type": "string",
                        "description": "Environment name used in table path (defaults to 'prod')"
                    },
                    "sr_ids": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of sound recording IDs to filter by"
                    },
                    "sr_version_ids": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of sound recording version IDs to filter by"
                    },
                    "service": {
                        "type": "string",
                        "description": "Service name to filter by"
                    },
                    "execution_type": {
                        "type": "array",
                        "items": {
                            "type": "string",
                            "enum": ["METADATA_UPDATE", "FULL_DELIVERY", "TAKEDOWN_DELIVERY"]
                        },
                        "description": "List of execution types to filter by. Valid values: 'METADATA_UPDATE', 'FULL_DELIVERY', 'TAKEDOWN_DELIVERY'"
                    },
                    "event_type": {
                        "type": "array",
                        "items": {
                            "type": "string",
                            "enum": ["success", "start", "not_eligible", "failure"]
                        },
                        "description": "List of statuses to filter by. Valid values: 'success', 'start', 'not_eligible', 'failure'"
                    }
                }
            }
        ),
        # Database Tools
        Tool(
            name="query_neo4j",
            description="Execute a Cypher query against Neo4j",
            inputSchema={
                "type": "object",
                "properties": {
                    "cypher": {
                        "type": "string",
                        "description": "The Cypher query to execute"
                    }
                },
                "required": ["cypher"]
            }
        ),
        Tool(
            name="find_neo4j_nodes",
            description="Find nodes in Neo4j by label and optional filters",
            inputSchema={
                "type": "object",
                "properties": {
                    "label": {
                        "type": "string",
                        "description": "Node label to search for"
                    },
                    "filters": {
                        "type": "object",
                        "description": "Optional property filters"
                    }
                },
                "required": ["label"]
            }
        ),
        Tool(
            name="query_snowflake",
            description="Execute a SQL query against Snowflake",
            inputSchema={
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "The SQL query to execute"
                    },
                    "params": {
                        "type": "array",
                        "description": "Optional query parameters"
                    }
                },
                "required": ["sql"]
            }
        ),
        Tool(
            name="get_snowflake_table_info",
            description="Get metadata about a Snowflake table",
            inputSchema={
                "type": "object",
                "properties": {
                    "table": {
                        "type": "string",
                        "description": "Table name"
                    }
                },
                "required": ["table"]
            }
        ),
        Tool(
            name="query_mysql",
            description="Execute a SQL query against MySQL",
            inputSchema={
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "The SQL query to execute"
                    },
                    "params": {
                        "type": "array",
                        "description": "Optional query parameters"
                    }
                },
                "required": ["sql"]
            }
        ),
        Tool(
            name="get_mysql_table_schema",
            description="Get schema information for a MySQL table",
            inputSchema={
                "type": "object",
                "properties": {
                    "table": {
                        "type": "string",
                        "description": "Table name"
                    }
                },
                "required": ["table"]
            }
        ),
        # AWS Tools
        Tool(
            name="invoke_lambda",
            description="Invoke a Lambda function",
            inputSchema={
                "type": "object",
                "properties": {
                    "function_name": {
                        "type": "string",
                        "description": "Name or ARN of the Lambda function"
                    },
                    "payload": {
                        "type": "object",
                        "description": "Input payload for the function"
                    },
                    "async_invoke": {
                        "type": "boolean",
                        "description": "Whether to invoke asynchronously",
                        "default": False
                    }
                },
                "required": ["function_name", "payload"]
            }
        ),
        Tool(
            name="list_lambda_functions",
            description="List all Lambda functions in the account",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        ),
        Tool(
            name="get_lambda_info",
            description="Get detailed information about a Lambda function",
            inputSchema={
                "type": "object",
                "properties": {
                    "function_name": {
                        "type": "string",
                        "description": "Name or ARN of the Lambda function"
                    }
                },
                "required": ["function_name"]
            }
        ),
        Tool(
            name="start_sfn_execution",
            description="Start a Step Functions execution",
            inputSchema={
                "type": "object",
                "properties": {
                    "state_machine_arn": {
                        "type": "string",
                        "description": "ARN of the state machine"
                    },
                    "input_data": {
                        "type": "object",
                        "description": "Input data for the execution"
                    },
                    "name": {
                        "type": "string",
                        "description": "Optional execution name"
                    }
                },
                "required": ["state_machine_arn", "input_data"]
            }
        ),
        Tool(
            name="describe_sfn_execution",
            description="Get details about a Step Functions execution",
            inputSchema={
                "type": "object",
                "properties": {
                    "execution_arn": {
                        "type": "string",
                        "description": "ARN of the execution"
                    }
                },
                "required": ["execution_arn"]
            }
        ),
        Tool(
            name="list_sfn_state_machines",
            description="List all Step Functions state machines",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        ),
        Tool(
            name="get_sfn_execution_history",
            description="Get execution history for a Step Functions execution",
            inputSchema={
                "type": "object",
                "properties": {
                    "execution_arn": {
                        "type": "string",
                        "description": "ARN of the execution"
                    }
                },
                "required": ["execution_arn"]
            }
        ),
        Tool(
            name="list_s3_buckets",
            description="List all S3 buckets",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        ),
        Tool(
            name="list_s3_objects",
            description="List objects in an S3 bucket",
            inputSchema={
                "type": "object",
                "properties": {
                    "bucket": {
                        "type": "string",
                        "description": "Bucket name"
                    },
                    "prefix": {
                        "type": "string",
                        "description": "Optional prefix to filter objects",
                        "default": ""
                    },
                    "max_keys": {
                        "type": "integer",
                        "description": "Maximum number of objects to return",
                        "default": 100
                    }
                },
                "required": ["bucket"]
            }
        ),
        Tool(
            name="get_s3_object",
            description="Get an object from S3",
            inputSchema={
                "type": "object",
                "properties": {
                    "bucket": {
                        "type": "string",
                        "description": "Bucket name"
                    },
                    "key": {
                        "type": "string",
                        "description": "Object key"
                    }
                },
                "required": ["bucket", "key"]
            }
        ),
        Tool(
            name="get_s3_bucket_size",
            description="Get total size of objects in an S3 bucket",
            inputSchema={
                "type": "object",
                "properties": {
                    "bucket": {
                        "type": "string",
                        "description": "Bucket name"
                    }
                },
                "required": ["bucket"]
            }
        ),
    ]


@server.call_tool()
async def call_tool(name: str, arguments: Dict[str, Any]) -> list[TextContent]:
    """Handle tool calls."""
    logger.info(f"Calling tool: {name} with arguments: {arguments}")

    try:
        # Database tools
        if name == "query_neo4j":
            result = neo4j_tools.query_neo4j(arguments["cypher"])
        elif name == "find_neo4j_nodes":
            result = neo4j_tools.find_neo4j_nodes(
                arguments["label"],
                arguments.get("filters")
            )
        elif name == "query_snowflake":
            result = snowflake_tools.query_snowflake(
                arguments["sql"],
                arguments.get("params")
            )
        elif name == "get_snowflake_table_info":
            result = snowflake_tools.get_snowflake_table_info(arguments["table"])
        elif name == "query_mysql":
            result = mysql_tools.query_mysql(
                arguments["sql"],
                arguments.get("params")
            )
        elif name == "get_mysql_table_schema":
            result = mysql_tools.get_mysql_table_schema(arguments["table"])

        # AWS tools
        elif name == "invoke_lambda":
            result = aws_tools.invoke_lambda(
                arguments["function_name"],
                arguments["payload"],
                arguments.get("async_invoke", False)
            )
        elif name == "list_lambda_functions":
            result = aws_tools.list_lambda_functions()
        elif name == "get_lambda_info":
            result = aws_tools.get_lambda_info(arguments["function_name"])
        elif name == "start_sfn_execution":
            result = aws_tools.start_sfn_execution(
                arguments["state_machine_arn"],
                arguments["input_data"],
                arguments.get("name")
            )
        elif name == "describe_sfn_execution":
            result = aws_tools.describe_sfn_execution(arguments["execution_arn"])
        elif name == "list_sfn_state_machines":
            result = aws_tools.list_sfn_state_machines()
        elif name == "get_sfn_execution_history":
            result = aws_tools.get_sfn_execution_history(arguments["execution_arn"])
        elif name == "list_s3_buckets":
            result = aws_tools.list_s3_buckets()
        elif name == "list_s3_objects":
            result = aws_tools.list_s3_objects(
                arguments["bucket"],
                arguments.get("prefix", ""),
                arguments.get("max_keys", 100)
            )
        elif name == "get_s3_object":
            result = aws_tools.get_s3_object(arguments["bucket"], arguments["key"])
        elif name == "get_s3_bucket_size":
            result = aws_tools.get_s3_bucket_size(arguments["bucket"])

        # Sound recording tools
        elif name == "get_delivery_history":
            result = sound_recording_tools.get_delivery_history(
                env=arguments.get("env"),
                sr_ids=arguments.get("sr_ids"),
                sr_version_ids=arguments.get("sr_version_ids"),
                service=arguments.get("service"),
                execution_type=arguments.get("execution_type"),
                event_type=arguments.get("event_type"),
            )
        else:
            result = {"error": f"Unknown tool: {name}"}

        return [TextContent(type="text", text=json.dumps(result, indent=2, default=_json_default))]

    except Exception as e:
        logger.error(f"Error calling tool {name}: {e}", exc_info=True)
        return [TextContent(type="text", text=json.dumps({
            "error": str(e),
            "tool": name
        }, indent=2))]


def main() -> None:
    """Run the MCP server."""
    args = _parse_args()
    logger.info("Sound Recording MCP Server started")

    if args.transport == "stdio":
        asyncio.run(_run_stdio_server())
        return

    import uvicorn

    app = _build_sse_app(sse_path=args.sse_path, message_path=args.message_path)
    logger.info(
        "Starting SSE transport on http://%s:%s%s (messages: %s)",
        args.host,
        args.port,
        _normalize_path(args.sse_path),
        _normalize_path(args.message_path),
    )
    uvicorn.run(app, host=args.host, port=args.port, log_level="info", access_log=False)


if __name__ == "__main__":
    main()