"""Agent integration commands: init, serve.

init  — generates a point-in-time tool manifest for AI agents
serve — local MCP proxy that injects auth for local AI clients,
        with optional multi-connection aggregation
"""

import asyncio
import json as json_mod
import logging
from dataclasses import dataclass
from pathlib import Path

import click
from aiohttp import web
from aiohttp import ClientSession as AiohttpSession

from mcpbundles.auth_resolve import resolve_auth, resolve_auth_for_named
from mcpbundles.config import AuthError, CLISettings
from mcpbundles.connections import ConnectionManager
from mcpbundles.manifest import generate_markdown_manifest, generate_json_manifest
from mcpbundles.mcp_client import hub_session, ConnectionError
from mcpbundles.output import render_success, render_error, render_warning, stdout_console

logger = logging.getLogger(__name__)


@click.command("init")
@click.option("--format", "fmt", type=click.Choice(["md", "json"]), default="md", help="Output format (default: md)")
@click.option("-o", "--output", "output_path", help="Output file path (default: mcpbundles_tools.md or .json)")
@click.option("--as", "connection", help="Use a named connection")
@click.pass_context
def init_cmd(ctx: click.Context, fmt: str, output_path: str | None, connection: str | None) -> None:
    """Generate a tool manifest snapshot for AI agents.

    Creates a point-in-time snapshot of available tools. AI agents can read
    this file to discover tools without making MCP calls.

    Examples:
        mcpbundles init                           # writes mcpbundles_tools.md
        mcpbundles init --format json             # writes mcpbundles_tools.json
        mcpbundles init -o tools.md               # custom output path
        mcpbundles init --as @hubspot             # from a named connection
    """
    asyncio.run(_init_impl(ctx, fmt, output_path, connection))


async def _init_impl(ctx: click.Context, fmt: str, output_path: str | None, connection: str | None) -> None:
    settings: CLISettings = ctx.obj["settings"]
    try:
        auth = resolve_auth(settings, connection)
    except (AuthError, SystemExit):
        raise SystemExit(1)

    try:
        async with hub_session(auth.token, auth.url, auth.endpoint) as session:
            result = await session.list_tools()
            tools = result.tools

            default_name = f"mcpbundles_tools.{fmt}"
            out = Path(output_path) if output_path else Path(default_name)

            if fmt == "json":
                generate_json_manifest(tools, out, connection_name=auth.conn_name)
            else:
                generate_markdown_manifest(tools, out, connection_name=auth.conn_name)

            render_success(f"Wrote {len(tools)} tools to {out}")
            if not output_path:
                stdout_console.print(f"  Add [cyan]{out.name}[/] to your project for AI agent discovery.")
                stdout_console.print("  [dim]This is a snapshot — run 'mcpbundles init' again to refresh.[/]")

    except (AuthError, ConnectionError) as exc:
        render_error(str(exc))
        raise SystemExit(1)


@click.command("serve")
@click.option("--port", default=8080, help="Local port to serve on (default: 8080)")
@click.option("--host", default="127.0.0.1", help="Local host to bind to (default: 127.0.0.1)")
@click.option("--as", "connection", help="Serve a specific named connection")
@click.option("--connections", "connection_names", help="Aggregate multiple connections (comma-separated)")
@click.pass_context
def serve_cmd(ctx: click.Context, port: int, host: str, connection: str | None, connection_names: str | None) -> None:
    """Start a local MCP proxy server.

    Exposes an unauthenticated local HTTP endpoint that proxies MCP
    requests to authenticated remote MCPBundles connections. Local AI
    clients connect here without needing credentials.

    With --connections, multiple named connections are aggregated into
    a single MCP surface. list_tools returns the combined tool set,
    and call_tool routes to the correct upstream automatically.

    Examples:
        mcpbundles serve                                        # serve Hub on localhost:8080
        mcpbundles serve --port 9090                            # custom port
        mcpbundles serve --as @hubspot                          # proxy a single connection
        mcpbundles serve --connections hub,hubspot,stripe       # aggregate multiple
    """
    asyncio.run(_serve_impl(ctx, port, host, connection, connection_names))


@dataclass
class _UpstreamTarget:
    name: str
    url: str
    endpoint: str
    auth: str


async def _serve_impl(ctx: click.Context, port: int, host: str, connection: str | None, connection_names: str | None) -> None:
    settings: CLISettings = ctx.obj["settings"]
    mgr = ConnectionManager()
    targets: list[_UpstreamTarget] = []

    if connection_names:
        for name in connection_names.split(","):
            name = name.strip().lstrip("@")
            conn = mgr.get(name)
            if not conn:
                render_error(f"Connection '{name}' not found. Use 'mcpbundles connect' first.")
                raise SystemExit(1)
            try:
                conn_auth = resolve_auth_for_named(conn, mgr, settings)
            except AuthError as exc:
                render_error(f"No auth for connection '{name}': {exc}")
                raise SystemExit(1)
            targets.append(_UpstreamTarget(name=conn.name, url=conn.url, endpoint=conn.endpoint, auth=conn_auth))
    else:
        try:
            auth = resolve_auth(settings, connection)
        except (AuthError, SystemExit):
            raise SystemExit(1)
        targets.append(_UpstreamTarget(name=auth.conn_name or "hub", url=auth.url, endpoint=auth.endpoint, auth=auth.token))

    if len(targets) == 1:
        await _serve_single(targets[0], host, port)
    else:
        await _serve_aggregated(targets, host, port)


def _build_auth_headers(auth: str) -> dict[str, str]:
    if auth.startswith("mb_"):
        return {"X-API-Key": auth}
    return {"Authorization": f"Bearer {auth}"}


async def _serve_single(target: _UpstreamTarget, host: str, port: int) -> None:

    async def handle_mcp(request: web.Request) -> web.Response:
        body = await request.read()
        target_url = f"{target.url}{target.endpoint}"
        headers = {"Content-Type": "application/json", **_build_auth_headers(target.auth)}

        for hdr in ("Mcp-Session-Id",):
            if hdr in request.headers:
                headers[hdr] = request.headers[hdr]

        async with AiohttpSession() as http_session:
            async with http_session.post(target_url, data=body, headers=headers) as resp:
                resp_body = await resp.read()
                return web.Response(
                    body=resp_body,
                    status=resp.status,
                    content_type=resp.content_type,
                    headers={
                        k: v for k, v in resp.headers.items()
                        if k.lower() in ("mcp-session-id", "content-type")
                    },
                )

    app = web.Application()
    app.router.add_post("/mcp", handle_mcp)
    app.router.add_get("/health", lambda _: web.json_response({"status": "ok", "connection": target.name}))

    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, host, port)
    await site.start()

    stdout_console.print(f"[bold]MCPBundles local proxy[/bold] serving on [cyan]http://{host}:{port}/mcp[/]")
    stdout_console.print(f"  Proxying → [dim]{target.url}{target.endpoint}[/]")
    stdout_console.print(f"  Connection: [cyan]{target.name}[/]")
    stdout_console.print()
    stdout_console.print("  AI clients can connect to this endpoint without authentication.")
    stdout_console.print("  Press Ctrl+C to stop.")

    try:
        await asyncio.Event().wait()
    except (KeyboardInterrupt, asyncio.CancelledError):
        pass
    finally:
        await runner.cleanup()
        stdout_console.print("\nProxy stopped.")


async def _serve_aggregated(targets: list[_UpstreamTarget], host: str, port: int) -> None:
    """Serve multiple upstream connections as a single unified MCP endpoint.

    list_tools merges tools from all upstreams. call_tool routes to the
    upstream that owns the requested tool. Tool ownership is discovered
    at startup and refreshed on each list_tools call.
    """
    tool_ownership: dict[str, _UpstreamTarget] = {}

    async def _discover_tools(http: AiohttpSession) -> list[dict]:
        """Fetch tools from all upstreams and build the ownership map."""
        all_tools: list[dict] = []
        tool_ownership.clear()

        for tgt in targets:
            url = f"{tgt.url}{tgt.endpoint}"
            headers = {"Content-Type": "application/json", **_build_auth_headers(tgt.auth)}
            init_payload = {"jsonrpc": "2.0", "method": "initialize", "id": 1, "params": {
                "protocolVersion": "2025-03-26",
                "capabilities": {},
                "clientInfo": {"name": "mcpbundles-serve", "version": "1.0.0"},
            }}
            try:
                async with http.post(url, json=init_payload, headers=headers) as resp:
                    session_id = resp.headers.get("Mcp-Session-Id")
                    if resp.status != 200:
                        logger.warning("Failed to initialize %s: HTTP %d", tgt.name, resp.status)
                        continue

                notify_payload = {"jsonrpc": "2.0", "method": "notifications/initialized"}
                notify_headers = {**headers}
                if session_id:
                    notify_headers["Mcp-Session-Id"] = session_id
                await http.post(url, json=notify_payload, headers=notify_headers)

                list_payload = {"jsonrpc": "2.0", "method": "tools/list", "id": 2, "params": {}}
                list_headers = {**headers}
                if session_id:
                    list_headers["Mcp-Session-Id"] = session_id
                async with http.post(url, json=list_payload, headers=list_headers) as list_resp:
                    if list_resp.status == 200:
                        data = await list_resp.json()
                        tools_result = data.get("result", {}).get("tools", [])
                        for tool in tools_result:
                            tool_name = tool.get("name", "")
                            if tool_name not in tool_ownership:
                                tool_ownership[tool_name] = tgt
                                all_tools.append(tool)
                            else:
                                logger.debug("Duplicate tool '%s' from %s (already owned by %s)",
                                             tool_name, tgt.name, tool_ownership[tool_name].name)
            except Exception as exc:
                logger.warning("Error discovering tools from %s: %s", tgt.name, exc)

        return all_tools

    async def handle_mcp(request: web.Request) -> web.Response:
        body_bytes = await request.read()
        try:
            body = json_mod.loads(body_bytes)
        except json_mod.JSONDecodeError:
            return web.json_response({"jsonrpc": "2.0", "error": {"code": -32700, "message": "Parse error"}}, status=400)

        method = body.get("method", "")
        req_id = body.get("id")

        async with AiohttpSession() as http:
            if method == "initialize":
                await _discover_tools(http)
                return web.json_response({
                    "jsonrpc": "2.0",
                    "id": req_id,
                    "result": {
                        "protocolVersion": "2025-03-26",
                        "capabilities": {"tools": {"listChanged": False}},
                        "serverInfo": {"name": "mcpbundles-aggregated-proxy", "version": "1.0.0"},
                    },
                })

            if method == "notifications/initialized":
                return web.Response(status=204)

            if method == "tools/list":
                tools_list = await _discover_tools(http)
                return web.json_response({
                    "jsonrpc": "2.0",
                    "id": req_id,
                    "result": {"tools": tools_list},
                })

            if method == "tools/call":
                tool_name = body.get("params", {}).get("name", "")
                owner = tool_ownership.get(tool_name)
                if not owner:
                    return web.json_response({
                        "jsonrpc": "2.0",
                        "id": req_id,
                        "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"},
                    })

                url = f"{owner.url}{owner.endpoint}"
                headers = {"Content-Type": "application/json", **_build_auth_headers(owner.auth)}
                init_payload = {"jsonrpc": "2.0", "method": "initialize", "id": 1, "params": {
                    "protocolVersion": "2025-03-26",
                    "capabilities": {},
                    "clientInfo": {"name": "mcpbundles-serve", "version": "1.0.0"},
                }}
                async with http.post(url, json=init_payload, headers=headers) as init_resp:
                    session_id = init_resp.headers.get("Mcp-Session-Id")

                notify_headers = {**headers}
                if session_id:
                    notify_headers["Mcp-Session-Id"] = session_id
                await http.post(url, json={"jsonrpc": "2.0", "method": "notifications/initialized"}, headers=notify_headers)

                call_headers = {**headers}
                if session_id:
                    call_headers["Mcp-Session-Id"] = session_id
                async with http.post(url, json=body, headers=call_headers) as call_resp:
                    resp_body = await call_resp.read()
                    return web.Response(
                        body=resp_body,
                        status=call_resp.status,
                        content_type=call_resp.content_type,
                    )

            return web.json_response({
                "jsonrpc": "2.0",
                "id": req_id,
                "error": {"code": -32601, "message": f"Method not supported: {method}"},
            })

    app = web.Application()
    app.router.add_post("/mcp", handle_mcp)
    conn_names = [t.name for t in targets]
    app.router.add_get("/health", lambda _: web.json_response({"status": "ok", "connections": conn_names}))

    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, host, port)
    await site.start()

    stdout_console.print(f"[bold]MCPBundles aggregated proxy[/bold] serving on [cyan]http://{host}:{port}/mcp[/]")
    stdout_console.print(f"  Aggregating {len(targets)} connections:")
    for t in targets:
        stdout_console.print(f"    [cyan]{t.name}[/] → [dim]{t.url}{t.endpoint}[/]")
    stdout_console.print()
    stdout_console.print("  list_tools returns merged tools from all connections.")
    stdout_console.print("  call_tool routes to the correct upstream automatically.")
    stdout_console.print("  Press Ctrl+C to stop.")

    try:
        await asyncio.Event().wait()
    except (KeyboardInterrupt, asyncio.CancelledError):
        pass
    finally:
        await runner.cleanup()
        stdout_console.print("\nProxy stopped.")
