"""Edge service orchestration app factory.

Builds a FastAPI application that:

* Exposes ``/healthz`` and ``/readyz`` SAGE health endpoints.
* **Reverse-proxies** all other requests to an externally-running LLM gateway
  process (sagellm-gateway or any OpenAI-compatible server) via HTTP.

The two ecosystems (SAGE and SageLLM) interact exclusively over HTTP.  There is
no in-process Python import of sagellm packages; sagellm-gateway runs as its own
process on a separate port (default: ``SagePorts.SAGELLM_GATEWAY`` = 8889).
"""

from __future__ import annotations

import httpx
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse

from sage.edge.core import EDGE_SERVICE_NAME, normalize_mount_path, probe_payload


def _attach_health_routes(app: FastAPI, llm_prefix: str | None, llm_upstream: str | None) -> None:
    """Attach /healthz and /readyz endpoints to the app (idempotent)."""
    existing_paths = {route.path for route in app.router.routes}
    llm_mounted = llm_upstream is not None

    if "/healthz" not in existing_paths:

        @app.get("/healthz", include_in_schema=False)
        async def healthz() -> dict[str, str | bool]:
            return probe_payload("ok", llm_mounted, llm_prefix)

    if "/readyz" not in existing_paths:

        @app.get("/readyz", include_in_schema=False)
        async def readyz() -> dict[str, str | bool]:
            return probe_payload("ready", llm_mounted, llm_prefix)


def _make_proxy_route(upstream_base: str, mount_path: str) -> None:
    """Return a wildcard proxy handler closure for *upstream_base*.

    Forwards every request under *mount_path* to *upstream_base* using httpx,
    streaming the response back to callers.  Works for both regular and
    server-sent-event (SSE) streaming responses.
    """
    strip_prefix = "" if mount_path == "/" else mount_path.rstrip("/")

    async def _proxy(request: Request) -> StreamingResponse:
        # Strip the edge mount prefix before forwarding
        path = request.url.path
        if strip_prefix and path.startswith(strip_prefix):
            path = path[len(strip_prefix):] or "/"
        query = request.url.query
        target = upstream_base.rstrip("/") + path
        if query:
            target = f"{target}?{query}"

        # Stream request body → upstream, stream response body → client
        async with httpx.AsyncClient(timeout=None) as client:
            upstream_req = client.build_request(
                method=request.method,
                url=target,
                headers={
                    k: v
                    for k, v in request.headers.items()
                    if k.lower() not in ("host", "content-length")
                },
                content=request.stream(),
            )
            upstream_resp = await client.send(upstream_req, stream=True)

        return StreamingResponse(
            upstream_resp.aiter_raw(),
            status_code=upstream_resp.status_code,
            headers=dict(upstream_resp.headers),
            media_type=upstream_resp.headers.get("content-type"),
            background=None,
        )

    return _proxy


def create_app(
    *,
    llm_upstream: str | None = None,
    llm_prefix: str | None = None,
) -> FastAPI:
    """Create the SAGE edge app.

    Args:
        llm_upstream: Base URL of the running LLM gateway process
            (e.g. ``"http://localhost:8889"``).  When ``None``, no proxy
            routes are added and only health endpoints are served.
        llm_prefix: Optional path prefix under which LLM routes are exposed
            on this edge app (default ``/``).

    Returns:
        A FastAPI application with health endpoints and, when *llm_upstream* is
        provided, wildcard reverse-proxy routes forwarding to the upstream.
    """
    mount_path = normalize_mount_path(llm_prefix)

    edge_app = FastAPI(title=EDGE_SERVICE_NAME, version="0.1.0")
    _attach_health_routes(edge_app, llm_prefix=llm_prefix, llm_upstream=llm_upstream)

    if llm_upstream:
        proxy_handler = _make_proxy_route(llm_upstream, mount_path)
        proxy_path = (mount_path.rstrip("/") or "") + "/{path:path}"
        edge_app.add_api_route(
            proxy_path,
            proxy_handler,
            methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"],
            include_in_schema=False,
        )
        edge_app.state.llm_upstream = llm_upstream
        edge_app.state.edge_mount_path = mount_path

    return edge_app
