"""FastAPI application with MCP server (Streamable HTTP transport)."""

from __future__ import annotations

import logging
import logging.config
from contextlib import asynccontextmanager

# Configure root logger so all cube_cloud.* loggers emit to stderr
logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s: %(message)s")

import uvicorn
from fastapi import FastAPI
from mcp.server.lowlevel import Server
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from mcp.types import Tool, TextContent
from starlette.middleware import Middleware
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Mount
from starlette.websockets import WebSocket, WebSocketDisconnect

from cube_common.apollo_client import ApolloClient, ApolloError
from cube_cloud.auth.middleware import APIKeyMiddleware
from cube_cloud.config import settings
from cube_cloud.tools import environments, entities, acr, clusters, kubectl_exec, apps
from cube_cloud._request_context import current_request

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# MCP Server definition
# ---------------------------------------------------------------------------

mcp_server = Server("cube-cloud")


@mcp_server.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="list_environments",
            description=(
                "List Apollo environments with their IDs and RIDs. "
                "Optionally filter by name substring."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "search": {
                        "type": "string",
                        "description": 'Optional name substring to filter by (case-insensitive). Example: "pep".',
                    },
                },
            },
        ),
        Tool(
            name="create_environment",
            description=(
                "Create a new Apollo environment and install the Apollo Control Plane module."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
                        "description": 'Environment name (e.g. "my-cube-test"). Must be unique.',
                    },
                    "accreditation": {
                        "type": "string",
                        "description": "Accreditation level. Defaults to DEV.",
                        "default": "DEV",
                    },
                },
                "required": ["name"],
            },
        ),
        Tool(
            name="replicate_environment",
            description=(
                "Replicate an Apollo environment into a new or existing environment. "
                "Copies modules, entities, config overrides, and reports secrets."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "source": {
                        "type": "string",
                        "description": "Source environment name or RID.",
                    },
                    "target": {
                        "type": "string",
                        "description": "Target environment name. Created if it doesn't exist.",
                    },
                    "accreditation": {
                        "type": "string",
                        "description": "Accreditation for new environments. Defaults to DEV.",
                        "default": "DEV",
                    },
                },
                "required": ["source", "target"],
            },
        ),
        Tool(
            name="install_entity",
            description="Install a Helm chart entity on an Apollo environment.",
            inputSchema={
                "type": "object",
                "properties": {
                    "environment": {
                        "type": "string",
                        "description": "Target environment name or RID.",
                    },
                    "entity_name": {
                        "type": "string",
                        "description": "Helm chart name (e.g. mosquitto).",
                    },
                    "product_id": {
                        "type": "string",
                        "description": 'Product maven coordinate "groupId:artifactId".',
                    },
                    "k8s_namespace": {
                        "type": "string",
                        "description": "Kubernetes namespace. Defaults to default.",
                        "default": "default",
                    },
                    "config_overrides": {
                        "type": "string",
                        "description": "Optional JSON config overrides.",
                    },
                },
                "required": ["environment", "entity_name", "product_id"],
            },
        ),
        Tool(
            name="entity_health",
            description="Get health and activity status of entities in an Apollo environment.",
            inputSchema={
                "type": "object",
                "properties": {
                    "environment": {
                        "type": "string",
                        "description": "Environment name or RID.",
                    },
                    "entity_id": {
                        "type": "string",
                        "description": "Optional entity display name to filter by.",
                    },
                },
                "required": ["environment"],
            },
        ),
        Tool(
            name="plan_details",
            description=(
                "Get detailed plan information for an entity, including tasks, events, and error logs."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "environment": {
                        "type": "string",
                        "description": "Environment name or RID.",
                    },
                    "entity_name": {
                        "type": "string",
                        "description": "Entity display name (substring match).",
                    },
                    "plan_index": {
                        "type": "integer",
                        "description": "Which plan to show (0 = most recent). Defaults to 0.",
                        "default": 0,
                    },
                },
                "required": ["environment", "entity_name"],
            },
        ),
        Tool(
            name="acr_get_token",
            description="Get an OAuth2 access token from Apollo Container Registry.",
            inputSchema={"type": "object", "properties": {}},
        ),
        Tool(
            name="apollo_publish_manifest",
            description="Publish an Apollo manifest YAML to make a Helm chart available for deployment.",
            inputSchema={
                "type": "object",
                "properties": {
                    "manifest_yaml": {
                        "type": "string",
                        "description": "The manifest YAML content to publish.",
                    },
                },
                "required": ["manifest_yaml"],
            },
        ),
        # Phase 2: Server-side Teleport + kubectl tools
        Tool(
            name="cube_status",
            description="Get the status of a Cube node (server-side kubectl via tbot).",
            inputSchema={
                "type": "object",
                "properties": {
                    "cluster": {
                        "type": "string",
                        "description": "Cube cluster name (from cube_list).",
                    },
                },
                "required": ["cluster"],
            },
        ),
        Tool(
            name="cube_list",
            description="List available Cube clusters accessible via server-side tbot.",
            inputSchema={"type": "object", "properties": {}},
        ),
        Tool(
            name="kubectl_exec",
            description=(
                "Execute a kubectl command server-side. "
                "Destructive commands (delete, apply, create, etc.) are blocked."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "cluster": {
                        "type": "string",
                        "description": "Target Cube cluster name (from cube_list).",
                    },
                    "command": {
                        "type": "string",
                        "description": 'kubectl command without the "kubectl" prefix. Example: "get pods -n default".',
                    },
                },
                "required": ["cluster", "command"],
            },
        ),
        Tool(
            name="app_list",
            description="List available Teleport apps (server-side via tbot).",
            inputSchema={
                "type": "object",
                "properties": {
                    "company": {
                        "type": "string",
                        "description": 'Optional filter by company label (e.g. "edgescaleai", "lear").',
                    },
                },
            },
        ),
        Tool(
            name="cube_cluster_login",
            description=(
                "Generate a self-contained kubeconfig for a Cube cluster. "
                "Returns a kubeconfig YAML with embedded TLS certs that works "
                "on any machine (no tsh or tbot needed locally)."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "cluster": {
                        "type": "string",
                        "description": "Cube cluster name (from cube_list).",
                    },
                },
                "required": ["cluster"],
            },
        ),
    ]


@mcp_server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    """Route tool calls to the appropriate handler.

    Credentials come from request.state (set by APIKeyMiddleware).
    Threaded via contextvars since MCP handlers don't receive the request.
    """
    req = current_request.get()
    client_id = req.state.apollo_client_id
    client_secret = req.state.apollo_client_secret

    async with ApolloClient(client_id, client_secret) as client:
        try:
            result = await _dispatch(client, name, arguments)
        except ApolloError as e:
            result = f"Error: {e}"

    return [TextContent(type="text", text=result)]


async def _dispatch(client: ApolloClient, name: str, args: dict) -> str:
    match name:
        case "list_environments":
            return await environments.list_environments(client, args.get("search"))
        case "create_environment":
            return await environments.create_environment(
                client, args["name"], args.get("accreditation", "DEV")
            )
        case "replicate_environment":
            return await environments.replicate_environment(
                client, args["source"], args["target"], args.get("accreditation", "DEV")
            )
        case "install_entity":
            return await environments.install_entity(
                client,
                args["environment"],
                args["entity_name"],
                args["product_id"],
                args.get("k8s_namespace", "default"),
                args.get("config_overrides"),
            )
        case "entity_health":
            return await entities.entity_health(
                client, args["environment"], args.get("entity_id")
            )
        case "plan_details":
            return await entities.plan_details(
                client, args["environment"], args["entity_name"], args.get("plan_index", 0)
            )
        case "acr_get_token":
            return await acr.acr_get_token(client)
        case "apollo_publish_manifest":
            return await acr.apollo_publish_manifest(client, args["manifest_yaml"])
        # Phase 2: these tools don't need the Apollo client
        case "cube_status" | "cube_list" | "kubectl_exec" | "app_list" | "cube_cluster_login":
            return await _dispatch_teleport(name, args)
        case _:
            return f"Unknown tool: {name}"


async def _dispatch_teleport(name: str, args: dict) -> str:
    """Dispatch tools that use server-side tbot/kubectl (no Apollo client needed)."""
    match name:
        case "cube_status":
            return await clusters.cube_status(args.get("cluster"))
        case "cube_list":
            return await clusters.cube_list()
        case "kubectl_exec":
            return await kubectl_exec.kubectl_exec(
                args["command"], cluster=args.get("cluster")
            )
        case "app_list":
            return await apps.app_list(args.get("company"))
        case "cube_cluster_login":
            return await clusters.cube_cluster_login(args["cluster"])
        case _:
            return f"Unknown teleport tool: {name}"


# ---------------------------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------------------------

app = FastAPI(title="cube-cloud", version="0.1.0")
app.add_middleware(APIKeyMiddleware)


@app.get("/health")
async def health():
    from cube_cloud.teleport.tbot_manager import tbot
    return {"status": "ok", "tbot": tbot.status(), "tunnel": _tunnel_conn_mgr.status()}


# ---------------------------------------------------------------------------
# Browser login flow (Cognito)
# ---------------------------------------------------------------------------

from pathlib import Path
from starlette.staticfiles import StaticFiles
from cube_cloud.auth.login_flow import auth_login, auth_authenticate, auth_session, setup_page

app.add_api_route("/auth/login", auth_login, methods=["GET"])
app.add_api_route("/auth/authenticate", auth_authenticate, methods=["POST"])
app.add_api_route("/auth/session/{session_id}", auth_session, methods=["GET"])
app.add_api_route("/setup", setup_page, methods=["GET"])
app.mount("/static", StaticFiles(directory=Path(__file__).parent / "static"), name="static")


# ---------------------------------------------------------------------------
# Tunnel relay (Phase 4) — WebSocket endpoint for app proxy
# ---------------------------------------------------------------------------

from cube_cloud.tunnel.connection_manager import ConnectionManager
from cube_cloud.tunnel.relay import TunnelRelay

_tunnel_conn_mgr = ConnectionManager()


@app.websocket("/tunnel")
async def tunnel_endpoint(ws: WebSocket):
    """WebSocket tunnel endpoint for cube-agent app proxy.

    The agent opens a WebSocket here, then sends CONNECT frames to
    request TCP connections to Teleport apps. Data is relayed
    bidirectionally via binary frames.
    """
    # Validate API key from query param or header
    api_key = ws.query_params.get("token") or ws.headers.get("authorization", "").removeprefix("Bearer ").strip()
    if not api_key:
        await ws.close(code=4001, reason="Missing authentication")
        return

    from cube_cloud.auth.api_keys import validate_key
    try:
        profile = await validate_key(api_key)
    except Exception as e:
        logger.error("Tunnel auth error: %s", e)
        await ws.close(code=4001, reason="Auth error")
        return
    if not profile:
        logger.warning("Tunnel auth failed: invalid API key")
        await ws.close(code=4001, reason="Invalid API key")
        return

    logger.info("Tunnel WebSocket authenticated, profile=%s", profile.get("profile"))
    await ws.accept()

    try:
        relay = TunnelRelay(ws, _tunnel_conn_mgr)
        await relay.run()
    except Exception as e:
        logger.error("Tunnel endpoint error: %s", e)


# ---------------------------------------------------------------------------
# Mount MCP at /mcp using StreamableHTTPSessionManager
# ---------------------------------------------------------------------------
# The session manager is an ASGI app that handles session lifecycle.
# We wrap it to inject the Starlette Request into a contextvar so MCP
# tool handlers can access auth state.

session_manager = StreamableHTTPSessionManager(
    app=mcp_server,
    json_response=True,
    stateless=True,
)


async def mcp_asgi_with_context(scope, receive, send):
    """ASGI wrapper that sets the current_request contextvar."""
    if scope["type"] == "http":
        request = Request(scope, receive, send)
        token = current_request.set(request)
        try:
            await session_manager.handle_request(scope, receive, send)
        finally:
            current_request.reset(token)
    else:
        await session_manager.handle_request(scope, receive, send)


app.mount("/mcp", mcp_asgi_with_context)


# ---------------------------------------------------------------------------
# Lifespan — start the MCP session manager's internal task group
# ---------------------------------------------------------------------------

@asynccontextmanager
async def lifespan(app: FastAPI):
    async with session_manager.run():
        logger.info("MCP session manager started")
        yield
    logger.info("MCP session manager stopped")


app.router.lifespan_context = lifespan


def cli():
    """Entry point for the cube-cloud CLI."""
    uvicorn.run(
        "cube_cloud.main:app",
        host=settings.host,
        port=settings.port,
        log_level=settings.log_level,
    )


if __name__ == "__main__":
    cli()
