"""Prompt Forge MCP server (stdio)."""
from __future__ import annotations

import asyncio
import json
import os
import sys
import time
from collections import defaultdict
from dataclasses import dataclass
from typing import Any
from urllib.parse import urlencode

import httpx
import mcp.server.stdio
import mcp.types as types
from mcp.server import Server

from promptforge_mcp.tools import load_tools


# Rate limiting
class RateLimiter:
    def __init__(self, max_requests_per_minute: int = 60):
        self.max_requests = max_requests_per_minute
        self.requests = defaultdict(list)

    def check_rate_limit(self, client_id: str = "default") -> bool:
        """Check if client is within rate limit. Returns True if allowed."""
        now = time.time()
        minute_ago = now - 60

        # Clean old requests
        self.requests[client_id] = [
            req_time for req_time in self.requests[client_id]
            if req_time > minute_ago
        ]

        # Check limit
        if len(self.requests[client_id]) >= self.max_requests:
            return False

        # Record this request
        self.requests[client_id].append(now)
        return True


rate_limiter = RateLimiter(
    max_requests_per_minute=int(os.getenv("MCP_RATE_LIMIT_PER_MINUTE", "60"))
)


@dataclass(frozen=True)
class MCPConfig:
    api_base: str
    gateway_url: str
    api_key: str | None
    admin_token: str | None
    org_id: str | None
    timeout: float

    @staticmethod
    def from_env() -> "MCPConfig":
        api_base = os.getenv("PROMPTFORGE_API_BASE", "https://api.secureprompt.tech")
        gateway_url = os.getenv("PROMPTFORGE_GATEWAY_URL", f"{api_base}/v1/gateway/requests")
        return MCPConfig(
            api_base=api_base,
            gateway_url=gateway_url,
            api_key=os.getenv("PROMPTFORGE_API_KEY"),
            admin_token=os.getenv("PROMPTFORGE_ADMIN_TOKEN"),
            org_id=os.getenv("PROMPTFORGE_ORG_ID"),
            timeout=float(os.getenv("PROMPTFORGE_TIMEOUT", "30")),
        )


TOOLS = load_tools()
TOOLS_BY_NAME = {tool["name"]: tool for tool in TOOLS}

server = Server("promptforge-mcp")

# L-8: Pooled httpx client instead of creating per-request
_http_client: httpx.AsyncClient | None = None


def _get_http_client() -> httpx.AsyncClient:
    """Get or create a pooled httpx client (L-8)."""
    global _http_client
    if _http_client is None:
        _http_client = httpx.AsyncClient(
            timeout=httpx.Timeout(30.0, connect=10.0),
            limits=httpx.Limits(max_connections=50, max_keepalive_connections=10),
        )
    return _http_client


@server.list_tools()
async def list_tools() -> list[types.Tool]:
    return [
        types.Tool(
            name=tool["name"],
            description=tool.get("description", ""),
            inputSchema=tool.get("input_schema") or tool.get("inputSchema") or {},
        )
        for tool in TOOLS
    ]


async def _request_json(
    method: str,
    url: str,
    *,
    headers: dict[str, str] | None = None,
    payload: dict[str, Any] | None = None,
    timeout: float = 30,
) -> dict[str, Any]:
    # L-8: Use pooled client instead of creating one per request
    client = _get_http_client()
    response = await client.request(method, url, json=payload, headers=headers, timeout=timeout)
    if response.status_code >= 400:
        raise ValueError(f"API request failed (HTTP {response.status_code})")
    return response.json()


def _resolve_org_id(config: MCPConfig, arguments: dict[str, Any]) -> str | None:
    return arguments.get("org_id") or config.org_id


@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> list[types.TextContent]:
    # Rate limiting check
    if not rate_limiter.check_rate_limit():
        raise ValueError("Rate limit exceeded. Please try again later.")

    config = MCPConfig.from_env()
    tool = TOOLS_BY_NAME.get(name)
    if not tool:
        raise ValueError(f"Unknown tool: {name}")

    if name == "promptforge.secure_generate":
        if not config.api_key:
            raise ValueError("PROMPTFORGE_API_KEY is required")
        headers = {"X-API-Key": config.api_key}
        result = await _request_json(
            "POST",
            config.gateway_url,
            headers=headers,
            payload=arguments,
            timeout=config.timeout,
        )
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.redact_and_tokenize":
        org_id = _resolve_org_id(config, arguments)
        if not config.admin_token or not org_id:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN and PROMPTFORGE_ORG_ID are required")
        payload = dict(arguments)
        payload.pop("org_id", None)
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/secure-mode/tokenize?org_id={org_id}"
        result = await _request_json("POST", url, headers=headers, payload=payload, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.policy_dry_run":
        org_id = _resolve_org_id(config, arguments)
        if not config.admin_token or not org_id:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN and PROMPTFORGE_ORG_ID are required")
        payload = dict(arguments)
        payload.pop("org_id", None)
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/policies/dry-run?org_id={org_id}"
        result = await _request_json("POST", url, headers=headers, payload=payload, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.audit_query":
        org_id = _resolve_org_id(config, arguments)
        if not config.admin_token or not org_id:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN and PROMPTFORGE_ORG_ID are required")
        params = {"org_id": org_id}
        for key in ["start_date", "end_date", "action", "user_id", "limit", "cursor"]:
            if key in arguments and arguments[key] is not None:
                params[key] = arguments[key]
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/audit-logs?{urlencode(params)}"
        result = await _request_json("GET", url, headers=headers, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.detokenize":
        org_id = _resolve_org_id(config, arguments)
        if not config.admin_token or not org_id:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN and PROMPTFORGE_ORG_ID are required")
        payload = dict(arguments)
        payload.pop("org_id", None)
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/secure-mode/detokenize?org_id={org_id}"
        result = await _request_json("POST", url, headers=headers, payload=payload, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.list_gateway_requests":
        org_id = _resolve_org_id(config, arguments)
        if not config.admin_token or not org_id:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN and PROMPTFORGE_ORG_ID are required")
        params: dict[str, Any] = {"org_id": org_id}
        for key in ["start_date", "end_date", "provider", "decision", "limit", "cursor"]:
            if key in arguments and arguments[key] is not None:
                params[key] = arguments[key]
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/gateway/requests?{urlencode(params)}"
        result = await _request_json("GET", url, headers=headers, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.list_secure_mode_entities":
        if not config.admin_token:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN is required")
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/secure-mode/entities"
        result = await _request_json("GET", url, headers=headers, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.get_secure_mode_config":
        org_id = _resolve_org_id(config, arguments)
        if not config.admin_token or not org_id:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN and PROMPTFORGE_ORG_ID are required")
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/secure-mode/config?org_id={org_id}"
        result = await _request_json("GET", url, headers=headers, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.get_billing_status":
        org_id = _resolve_org_id(config, arguments)
        if not config.admin_token or not org_id:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN and PROMPTFORGE_ORG_ID are required")
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/billing/status?org_id={org_id}"
        result = await _request_json("GET", url, headers=headers, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.list_billing_plans":
        if not config.admin_token:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN is required")
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/billing/plans"
        result = await _request_json("GET", url, headers=headers, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.list_providers":
        org_id = _resolve_org_id(config, arguments)
        if not config.admin_token or not org_id:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN and PROMPTFORGE_ORG_ID are required")
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/providers?org_id={org_id}"
        result = await _request_json("GET", url, headers=headers, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    if name == "promptforge.test_provider":
        if not config.admin_token:
            raise ValueError("PROMPTFORGE_ADMIN_TOKEN is required")
        provider_id = arguments.get("provider_id")
        if not provider_id:
            raise ValueError("provider_id is required")
        headers = {"Authorization": f"Bearer {config.admin_token}"}
        url = f"{config.api_base}/v1/providers/{provider_id}/test"
        result = await _request_json("POST", url, headers=headers, payload={}, timeout=config.timeout)
        return [types.TextContent(type="text", text=json.dumps(result))]

    raise ValueError(f"Unhandled tool: {name}")


async def run_stdio() -> None:
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            server.create_initialization_options(),
        )


async def run_sse(host: str = "0.0.0.0", port: int = 8765) -> None:
    import uvicorn
    from mcp.server.sse import SseServerTransport

    transport = SseServerTransport("/messages")

    async def app(scope, receive, send):
        if scope["type"] != "http":
            return
        path = scope.get("path", "")
        if path == "/sse":
            async with transport.connect_sse(scope, receive, send) as (read_stream, write_stream):
                await server.run(
                    read_stream,
                    write_stream,
                    server.create_initialization_options(),
                )
        elif path == "/messages":
            await transport.handle_post_message(scope, receive, send)
        else:
            await send({"type": "http.response.start", "status": 404, "headers": [(b"content-type", b"text/plain")]})
            await send({"type": "http.response.body", "body": b"Not Found"})

    max_connections = int(os.getenv("MCP_MAX_CONNECTIONS", "100"))
    config = uvicorn.Config(
        app,
        host=host,
        port=port,
        log_level="info",
        limit_concurrency=max_connections,
        timeout_keep_alive=30,
    )
    srv = uvicorn.Server(config)
    await srv.serve()


def main() -> None:
    mode = os.getenv("MCP_TRANSPORT", "stdio")
    if "--http" in sys.argv:
        mode = "http"
    host = os.getenv("MCP_HOST", "0.0.0.0")
    port = int(os.getenv("MCP_PORT", "8765"))

    if mode == "http":
        asyncio.run(run_sse(host=host, port=port))
    else:
        asyncio.run(run_stdio())


if __name__ == "__main__":
    main()
