from __future__ import annotations

import json
from collections.abc import Awaitable, Callable
from typing import Any

import anyio
import uvicorn
from mcp.server.fastmcp import FastMCP

from .backend_client import BackendClient
from .cloud_auth import CloudAuthBroker
from .cloud_context import current_cloud_context
from .config import get_cloud_http_allowed_hosts, get_cloud_http_host, get_cloud_http_port
from .errors import QingflowApiError, is_auth_like_error
from .server_app_builder import build_builder_server
from .server_app_user import build_user_server


BUILDER_TOOL_PREFIX = "builder_"


ASGIApp = Callable[
    [
        dict[str, Any],
        Callable[[], Awaitable[dict[str, Any]]],
        Callable[[dict[str, Any]], Awaitable[None]],
    ],
    Awaitable[None],
]


class CloudMcpAuthApp:
    def __init__(
        self,
        app: ASGIApp,
        broker: CloudAuthBroker,
        *,
        allowed_hosts: list[str] | None = None,
        internal_host: str | None = None,
    ) -> None:
        self._app = app
        self._broker = broker
        self._allowed_hosts = [
            host.strip().lower()
            for host in (allowed_hosts if allowed_hosts is not None else get_cloud_http_allowed_hosts())
            if host.strip()
        ]
        self._internal_host = internal_host or f"127.0.0.1:{get_cloud_http_port()}"

    async def __call__(
        self,
        scope: dict[str, Any],
        receive: Callable[[], Awaitable[dict[str, Any]]],
        send: Callable[[dict[str, Any]], Awaitable[None]],
    ) -> None:
        if scope.get("type") != "http":
            await self._app(scope, receive, send)
            return

        host = _extract_host(scope)
        if not _is_allowed_host(host, self._allowed_hosts):
            await self._plain_response(send, 421, "Invalid Host header")
            return
        scope = _rewrite_host(scope, self._internal_host)

        path = str(scope.get("path") or "")
        if path in {"/health", "/healthz", "/readyz"}:
            await self._json_response(send, 200, {"status": "ok"})
            return

        credential = self._extract_credential(scope)
        if credential is None:
            await self._json_response(
                send,
                401,
                {
                    "error": "missing_mcp_credential",
                    "message": "Authorization: Bearer <mcp_credential> is required",
                },
            )
            return

        try:
            cloud_context = await anyio.to_thread.run_sync(self._broker.resolve, credential)
        except QingflowApiError as error:
            status = 401 if is_auth_like_error(error) else 500
            await self._json_response(
                send,
                status,
                {
                    "error": error.category,
                    "message": error.message,
                    "request_id": error.request_id,
                    "backend_code": error.backend_code,
                },
            )
            return

        token = current_cloud_context.set(cloud_context)
        try:
            scope = _ensure_streamable_http_accept(scope)
            await self._app(scope, receive, send)
        finally:
            current_cloud_context.reset(token)

    def _extract_credential(self, scope: dict[str, Any]) -> str | None:
        headers = {
            key.decode("latin1").lower(): value.decode("latin1")
            for key, value in scope.get("headers", [])
        }
        authorization = headers.get("authorization", "")
        if authorization.lower().startswith("bearer "):
            credential = authorization[7:].strip()
            return credential or None
        credential = headers.get("x-qingflow-client-id", "").strip()
        return credential or None

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

    async def _plain_response(
        self,
        send: Callable[[dict[str, Any]], Awaitable[None]],
        status: int,
        text: str,
    ) -> None:
        body = text.encode("utf-8")
        await send(
            {
                "type": "http.response.start",
                "status": status,
                "headers": [
                    (b"content-type", b"text/plain; charset=utf-8"),
                    (b"content-length", str(len(body)).encode("ascii")),
                ],
            }
        )
        await send({"type": "http.response.body", "body": body})


class LazyCloudApp:
    def __init__(self) -> None:
        self._app: CloudMcpAuthApp | None = None

    async def __call__(
        self,
        scope: dict[str, Any],
        receive: Callable[[], Awaitable[dict[str, Any]]],
        send: Callable[[dict[str, Any]], Awaitable[None]],
    ) -> None:
        if self._app is None:
            self._app = build_cloud_app()
        await self._app(scope, receive, send)


def build_cloud_user_app() -> CloudMcpAuthApp:
    return build_cloud_app()


def build_cloud_app() -> CloudMcpAuthApp:
    server = build_combined_server()
    broker = CloudAuthBroker.from_config(backend=BackendClient())
    return CloudMcpAuthApp(server.streamable_http_app(), broker)


def build_combined_server() -> FastMCP:
    server = build_user_server()
    builder_server = build_builder_server()
    server._mcp_server.name = "Qingflow Combined MCP"
    _register_prefixed_builder_tools(server, builder_server)
    _set_combined_instructions(server, builder_server)
    return server


def _register_prefixed_builder_tools(target_server: FastMCP, builder_server: FastMCP) -> None:
    for original_name, tool in builder_server._tool_manager._tools.items():
        target_server._tool_manager.add_tool(
            tool.fn,
            name=f"{BUILDER_TOOL_PREFIX}{original_name}",
            title=tool.title,
            description=tool.description,
            annotations=tool.annotations,
            icons=tool.icons,
            meta=tool.meta,
        )


def _set_combined_instructions(user_server: FastMCP, builder_server: FastMCP) -> None:
    user_instructions = user_server.instructions or ""
    builder_instructions = builder_server.instructions or ""
    user_server._mcp_server.instructions = f"""Use this combined Qingflow MCP server for both user operational workflows and builder workflows.

## Tool Namespaces

User tools keep their original names, such as `app_get`, `record_list`, `task_list`, `portal_get`,
and `chart_get`.
Builder tools are exposed with the `{BUILDER_TOOL_PREFIX}` prefix. For example, Builder `app_get` is
`builder_app_get`, Builder `view_get` is `builder_view_get`, and Builder `builder_tool_contract` is
`builder_builder_tool_contract`.
When reading the Builder workflow instructions below, add `{BUILDER_TOOL_PREFIX}` to every Builder tool
name before calling it.

## User Workflows

{user_instructions}

## Builder Workflows

{builder_instructions}
"""


def _extract_host(scope: dict[str, Any]) -> str:
    for key, value in scope.get("headers", []):
        if key.lower() == b"host":
            return value.decode("latin1").strip().lower()
    server = scope.get("server")
    if isinstance(server, tuple) and len(server) >= 2:
        return f"{server[0]}:{server[1]}".lower()
    return ""


def _is_allowed_host(host: str, allowed_hosts: list[str]) -> bool:
    normalized_host = host.strip().lower().rstrip(".")
    host_without_port = _strip_port(normalized_host)
    for allowed in allowed_hosts:
        normalized_allowed = allowed.strip().lower().rstrip(".")
        if normalized_allowed == "*":
            return True
        if normalized_allowed.startswith("*."):
            suffix = normalized_allowed[1:]
            if host_without_port.endswith(suffix) and host_without_port != normalized_allowed[2:]:
                return True
            continue
        if normalized_host == normalized_allowed or host_without_port == normalized_allowed:
            return True
    return False


def _strip_port(host: str) -> str:
    if host.startswith("["):
        end = host.find("]")
        return host[1:end] if end > 0 else host
    if host.count(":") == 1:
        return host.split(":", 1)[0]
    return host


def _rewrite_host(scope: dict[str, Any], internal_host: str) -> dict[str, Any]:
    rewritten = dict(scope)
    headers = []
    host_seen = False
    for key, value in scope.get("headers", []):
        if key.lower() == b"host":
            headers.append((key, internal_host.encode("latin1")))
            host_seen = True
        else:
            headers.append((key, value))
    if not host_seen:
        headers.append((b"host", internal_host.encode("latin1")))
    rewritten["headers"] = headers
    return rewritten


def _ensure_streamable_http_accept(scope: dict[str, Any]) -> dict[str, Any]:
    existing_values: list[str] = []
    headers_without_accept: list[tuple[bytes, bytes]] = []
    for key, value in scope.get("headers", []):
        if key.lower() == b"accept":
            existing_values.append(value.decode("latin1"))
        else:
            headers_without_accept.append((key, value))

    required = ["application/json", "text/event-stream"]
    merged_values = [value.strip() for value in ",".join(existing_values).split(",") if value.strip()]
    normalized_media_types = {
        value.split(";", 1)[0].strip().lower()
        for value in merged_values
    }
    if not merged_values or "*/*" in normalized_media_types:
        merged_values = list(required)
    else:
        for media_type in required:
            if media_type not in normalized_media_types:
                merged_values.append(media_type)

    rewritten = dict(scope)
    rewritten["headers"] = headers_without_accept + [
        (b"accept", ", ".join(merged_values).encode("latin1"))
    ]
    return rewritten


app = LazyCloudApp()


def main() -> None:
    uvicorn.run(build_cloud_app(), host=get_cloud_http_host(), port=get_cloud_http_port())


if __name__ == "__main__":
    main()
