import os
import json
import logging
import re
import urllib.parse
from typing import Dict, Any, List, Optional
import httpx
from dotenv import load_dotenv

from mcp.server import Server
from mcp.types import Tool, TextContent
from mcp.server.sse import SseServerTransport
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, Response

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("openapi-mcp-server")

# Load environment
env_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".env")
load_dotenv(dotenv_path=env_path)

TRANSPORT = os.getenv("TRANSPORT", "stdio").lower()
PORT = int(os.getenv("PORT", "3000"))
FALLBACK_API_KEY = os.getenv("EGOI_API_KEY", "").strip()
API_BASE_URL = os.getenv("API_BASE_URL", "https://api.egoiapp.com").rstrip("/")
EXPOSE_TAGS_RAW = os.getenv("EXPOSE_TAGS", "")
EXPOSE_TAGS = [t.strip() for t in EXPOSE_TAGS_RAW.split(",") if t.strip()] if EXPOSE_TAGS_RAW else None

server = Server("openapi-mcp-server")

tools_list = []
tool_to_endpoint = {}
api_spec = {}

def build_tools_mapping(spec: dict):
    paths = spec.get("paths", {})
    for path_pattern, path_item in paths.items():
        for method, operation in path_item.items():
            http_method = method.lower()
            if http_method not in ["get", "post", "put", "delete", "patch"]:
                continue
                
            # Tag filtering
            if EXPOSE_TAGS:
                op_tags = operation.get("tags", [])
                if not any(tag in EXPOSE_TAGS for tag in op_tags):
                    continue
            
            # Name conversion: operationId camelCase to snake_case
            operation_id = operation.get("operationId", "")
            if operation_id:
                tool_name = re.sub(r'([A-Z]+)', r'_\1', operation_id).lower().replace("__", "_").strip("_")
                tool_name = re.sub(r'[^a-z0-9_-]', '_', tool_name)
            else:
                tool_name = f"{http_method}_{path_pattern}"
                tool_name = re.sub(r'[^a-z0-9_-]', '_', tool_name.lower())
                tool_name = re.sub(r'_+', '_', tool_name).strip("_")

            # Ensure uniqueness and limit length to 64
            tool_name = tool_name[:64]
            suffix = 1
            while tool_name in tool_to_endpoint:
                tool_name = f"{tool_name[:60]}_{suffix}"[:64]
                suffix += 1

            summary = operation.get("summary", "")
            description = operation.get("description", "")
            tags_str = f"Tags: [{', '.join(operation.get('tags', []))}]" if operation.get("tags") else ""
            
            full_description = "\n\n".join(filter(None, [
                summary,
                description,
                f"HTTP Method: {http_method.upper()}",
                f"Path: {path_pattern}",
                tags_str,
                "Authentication: Pass your E-goi API Key in the 'apiKey' parameter. If not provided, it falls back to the server's EGOI_API_KEY environment variable."
            ]))

            path_params = []
            query_params = []
            header_params = []
            input_schema_properties = {}
            required_fields = []

            # Process parameters
            for param in operation.get("parameters", []):
                name = param.get("name")
                schema = param.get("schema", {"type": "string"})
                param_desc = param.get("description", "")
                
                mcp_prop = {
                    "type": schema.get("type", "string"),
                    "description": param_desc
                }
                if "enum" in schema:
                    mcp_prop["enum"] = schema["enum"]
                if "default" in schema:
                    mcp_prop["default"] = schema["default"]

                input_schema_properties[name] = mcp_prop
                if param.get("required"):
                    required_fields.append(name)

                param_in = param.get("in")
                if param_in == "path":
                    path_params.append({"name": name, "schema": schema})
                elif param_in == "query":
                    query_params.append({"name": name, "schema": schema})
                elif param_in == "header":
                    header_params.append({"name": name, "schema": schema})

            # Process Request Body (application/json)
            body_schema = None
            req_body = operation.get("requestBody")
            if req_body and "content" in req_body:
                json_content = req_body["content"].get("application/json")
                if json_content and "schema" in json_content:
                    body_schema = json_content["schema"]
                    if body_schema.get("type") == "object" and "properties" in body_schema:
                        for prop_name, prop_schema in body_schema["properties"].items():
                            mcp_prop_name = prop_name
                            # Prevent collision with path/query parameters
                            if mcp_prop_name in input_schema_properties:
                                mcp_prop_name = f"body_{prop_name}"
                            
                            mcp_prop = {
                                "type": prop_schema.get("type", "string"),
                                "description": prop_schema.get("description", "")
                            }
                            if "enum" in prop_schema:
                                mcp_prop["enum"] = prop_schema["enum"]
                            if "properties" in prop_schema:
                                mcp_prop["properties"] = prop_schema["properties"]
                            if "items" in prop_schema:
                                mcp_prop["items"] = prop_schema["items"]

                            input_schema_properties[mcp_prop_name] = mcp_prop
                            if "required" in body_schema and prop_name in body_schema["required"]:
                                required_fields.append(mcp_prop_name)
                    else:
                        mcp_prop = {
                            "type": body_schema.get("type", "object"),
                            "description": body_schema.get("description", "The JSON payload request body.")
                        }
                        if "items" in body_schema:
                            mcp_prop["items"] = body_schema["items"]
                        if "properties" in body_schema:
                            mcp_prop["properties"] = body_schema["properties"]
                        input_schema_properties["body"] = mcp_prop
                        if req_body.get("required"):
                            required_fields.append("body")

            # Inject the dynamic apiKey parameter
            input_schema_properties["apiKey"] = {
                "type": "string",
                "description": "Your E-goi API Key. Takes precedence over EGOI_API_KEY environment variable."
            }

            input_schema = {
                "type": "object",
                "properties": input_schema_properties,
                "required": required_fields
            }

            tools_list.append(Tool(
                name=tool_name,
                description=full_description,
                inputSchema=input_schema
            ))

            tool_to_endpoint[tool_name] = {
                "path": path_pattern,
                "method": http_method,
                "pathParams": path_params,
                "queryParams": query_params,
                "headerParams": header_params,
                "bodySchema": body_schema
            }

# Load OpenAPI schema
local_spec_path_same = os.path.join(os.path.dirname(os.path.abspath(__file__)), "openapi-resolved.json")
local_spec_path_parent = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "openapi-resolved.json")
local_spec_path = local_spec_path_same if os.path.exists(local_spec_path_same) else local_spec_path_parent

if os.path.exists(local_spec_path):
    logger.info(f"Loading pre-resolved OpenAPI spec from local file: {local_spec_path}...")
    with open(local_spec_path, "r", encoding="utf-8") as f:
        api_spec = json.load(f)
        build_tools_mapping(api_spec)
        logger.info(f"Server initialized. Registered {len(tools_list)} tools.")
else:
    logger.error("openapi-resolved.json not found in parent directory. Run the resolve_openapi script first.")
    raise FileNotFoundError("openapi-resolved.json not found.")

@server.list_tools()
async def list_tools() -> List[Tool]:
    return tools_list

@server.call_tool()
async def call_tool(name: str, arguments: Optional[Dict[str, Any]]) -> List[TextContent]:
    endpoint = tool_to_endpoint.get(name)
    if not endpoint:
        raise ValueError(f"Tool not found: {name}")

    args = arguments or {}

    # Resolve API Key
    api_key = (args.get("apiKey") or FALLBACK_API_KEY or "").strip()
    if not api_key:
        return [TextContent(
            type="text",
            text="Error: An API Key is required to perform this action. Please provide it using the 'apiKey' parameter in your tool call, or set the EGOI_API_KEY environment variable on the server."
        )]

    # 1. Resolve path parameters
    resolved_path = endpoint["path"]
    for param in endpoint["pathParams"]:
        p_name = param["name"]
        if p_name not in args:
            return [TextContent(
                type="text",
                text=f"Error: Missing required path parameter '{p_name}'"
            )]
        val = args[p_name]
        resolved_path = resolved_path.replace(f"{{{p_name}}}", urllib.parse.quote(str(val)))

    # 2. Resolve query parameters
    resolved_query_params = {}
    for param in endpoint["queryParams"]:
        p_name = param["name"]
        if p_name in args:
            resolved_query_params[p_name] = args[p_name]

    # 3. Resolve request body
    resolved_body = None
    body_schema = endpoint["bodySchema"]
    if body_schema:
        if body_schema.get("type") == "object" and "properties" in body_schema:
            resolved_body = {}
            for prop_name in body_schema["properties"].keys():
                mcp_prop_name = prop_name
                # check collisions
                if any(p["name"] == prop_name for p in endpoint["pathParams"]) or \
                   any(q["name"] == prop_name for q in endpoint["queryParams"]):
                    mcp_prop_name = f"body_{prop_name}"
                
                if mcp_prop_name in args:
                    resolved_body[prop_name] = args[mcp_prop_name]
        else:
            resolved_body = args.get("body")

    # 4. Resolve headers
    resolved_headers = {
        "Apikey": api_key,
        "accept": "application/json"
    }
    if endpoint["method"] not in ["get", "delete"]:
        resolved_headers["Content-Type"] = "application/json"

    for param in endpoint["headerParams"]:
        p_name = param["name"]
        if p_name in args:
            resolved_headers[p_name] = str(args[p_name])

    # 5. Build URL
    servers = api_spec.get("servers", [])
    spec_server_url = servers[0].get("url", "https://api.egoiapp.com") if servers else "https://api.egoiapp.com"
    api_base_url = os.getenv("API_BASE_URL", spec_server_url).rstrip("/")
    full_url = f"{api_base_url}{resolved_path}"

    is_write = endpoint["method"] in ["post", "put", "patch"]
    timeout = 60.0 if is_write else 10.0

    logger.info(f"Executing {endpoint['method'].upper()} {full_url}...")

    async with httpx.AsyncClient(timeout=timeout) as client:
        try:
            response = await client.request(
                method=endpoint["method"],
                url=full_url,
                params=resolved_query_params,
                json=resolved_body,
                headers=resolved_headers
            )
            
            # Format output
            try:
                res_data = response.json()
                text_out = json.dumps(res_data, indent=2)
            except Exception:
                text_out = response.text

            return [TextContent(type="text", text=text_out)]
        except httpx.HTTPStatusError as err:
            logger.error(f"API Error: {err}")
            try:
                err_data = err.response.json()
                err_txt = json.dumps(err_data, indent=2)
            except Exception:
                err_txt = err.response.text
            return [TextContent(type="text", text=f"API Error {err.response.status_code}: {err_txt}")]
        except Exception as err:
            logger.error(f"Request failed: {err}")
            return [TextContent(type="text", text=f"Request Error: {str(err)}")]

app = FastAPI(title="E-goi OpenAPI MCP Server")
sse = SseServerTransport("/messages")

@app.get("/")
async def root():
    return HTMLResponse(
        content="<html><body><h1>E-goi OpenAPI MCP Server (Python)</h1>"
                "<p>Ready and listening for SSE connections at <code>/sse</code></p></body></html>"
    )

@app.get("/sse")
async def sse_endpoint(request: Request):
    async with sse.connect_sse(
        request.scope, request.receive, request._send
    ) as streams:
        await server.run(
            streams[0],
            streams[1],
            server.create_initialization_options()
        )
    return Response()

@app.post("/messages")
async def messages_endpoint(request: Request):
    await sse.handle_post_message(request.scope, request.receive, request._send)

if __name__ == "__main__":
    if TRANSPORT == "sse":
        logger.info(f"Starting FastAPI SSE server on port {PORT}...")
        import uvicorn
        uvicorn.run(app, host="0.0.0.0", port=PORT)
    else:
        logger.info("Starting Stdio server...")
        import asyncio
        from mcp.server.stdio import stdio_server
        
        async def run_stdio():
            async with stdio_server() as (read_stream, write_stream):
                await server.run(
                    read_stream,
                    write_stream,
                    server.create_initialization_options()
                )

        asyncio.run(run_stdio())