"""FastAPI reverse proxy for multiple Streamlit applications."""

import asyncio
import contextlib
from pathlib import Path
from typing import Any

import httpx
import websockets
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.templating import Jinja2Templates

from proxy.app_discovery import discover_apps

app = FastAPI(title="Fansifter Prototypes")
templates = Jinja2Templates(directory=str(Path(__file__).parent / "templates"))

# Determine apps directory - try local first, then Docker path
apps_dir = Path(__file__).parent.parent / "apps"
if not apps_dir.exists():
    apps_dir = Path("/app/apps")

# Discover apps once at startup and store in memory
APP_MAP = discover_apps(apps_dir)

# HTTP client for proxying requests
http_client = httpx.AsyncClient(timeout=30.0)


@app.on_event("startup")
async def startup_event() -> None:
    """Log discovered apps on startup."""
    print(f"Discovered {len(APP_MAP)} apps:")
    for name, metadata in APP_MAP.items():
        print(
            f"  - {metadata['emoji']} {metadata['display_name']} ({name}) -> http://localhost:{metadata['port']}"
        )


@app.on_event("shutdown")
async def shutdown_event() -> None:
    """Close HTTP client on shutdown."""
    await http_client.aclose()


@app.get("/", response_class=HTMLResponse)
async def read_root(request: Request) -> HTMLResponse:
    """Serve the navigation page with grouped demo/example apps via Jinja template."""
    demo_apps = {
        name: meta
        for name, meta in APP_MAP.items()
        if (meta.get("semantic_type") or meta.get("type", "demo") or meta.get("runtime")) == "demo"
    }
    example_apps = {
        name: meta
        for name, meta in APP_MAP.items()
        if (meta.get("semantic_type") or meta.get("type") or meta.get("runtime")) == "example"
    }
    return templates.TemplateResponse(
        "index.html",
        {
            "request": request,
            "demo_apps": demo_apps,
            "example_apps": example_apps,
        },
    )


@app.get("/hello/")
async def hello() -> dict[str, Any]:
    return {"status": "ok"}


@app.websocket("/{app_name}/{path:path}")
async def websocket_proxy(websocket: WebSocket, app_name: str, path: str) -> None:
    """
    Proxy WebSocket connections to apps.

    Handles all WebSocket paths for any app (Streamlit, FastAPI, etc.).
    """
    if app_name not in APP_MAP:
        await websocket.close(code=1008)  # Policy Violation
        return

    target_port = APP_MAP[app_name]["port"]
    target_url = f"ws://localhost:{target_port}/{path}"

    await websocket.accept()

    client_closed = False
    server_closed = False

    # Forward cookies and other headers to the backend
    additional_headers = []
    if "cookie" in websocket.headers:
        additional_headers.append(("cookie", websocket.headers["cookie"]))
    if "user-agent" in websocket.headers:
        additional_headers.append(("user-agent", websocket.headers["user-agent"]))

    try:
        async with websockets.connect(
            target_url, additional_headers=additional_headers
        ) as remote_ws:  # type: ignore
            # Create tasks for bidirectional communication
            async def forward_to_remote() -> None:
                """Forward messages from client to Streamlit app."""
                nonlocal client_closed
                try:
                    while True:
                        message = await websocket.receive()
                        if message.get("type") == "websocket.disconnect":
                            client_closed = True
                            break
                        if "text" in message:
                            data = message["text"]
                            await remote_ws.send(data)  # type: ignore
                        elif "bytes" in message:
                            data = message["bytes"]
                            await remote_ws.send(data)  # type: ignore
                except WebSocketDisconnect:
                    client_closed = True
                except Exception:
                    client_closed = True

            async def forward_to_client() -> None:
                """Forward messages from Streamlit app to client."""
                nonlocal server_closed, client_closed
                try:
                    async for message in remote_ws:  # type: ignore
                        if client_closed:
                            break
                        try:
                            if isinstance(message, str):
                                await websocket.send_text(message)
                            else:
                                await websocket.send_bytes(message)
                        except RuntimeError as e:
                            if "after sending 'websocket.close'" in str(e):
                                client_closed = True
                                break
                            raise
                except Exception:
                    pass
                finally:
                    server_closed = True

            # Run both directions concurrently
            await asyncio.gather(
                forward_to_remote(),
                forward_to_client(),
                return_exceptions=True,
            )

    except Exception:
        pass
    finally:
        if not client_closed:
            with contextlib.suppress(RuntimeError):
                await websocket.close()


@app.api_route(
    "/{app_name}/{path:path}",
    methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"],
)
async def http_proxy(request: Request, app_name: str, path: str) -> StreamingResponse:
    """
    Proxy HTTP requests to Streamlit apps.

    Handles all HTTP methods and streams responses back to the client.
    """
    if app_name not in APP_MAP:
        return StreamingResponse(
            content=iter([b"App not found"]),
            status_code=404,
            media_type="text/plain",
        )

    target_port = APP_MAP[app_name]["port"]
    target_url = f"http://localhost:{target_port}/{path}"

    # Build the proxied request
    headers = dict(request.headers)
    # Remove host header to avoid conflicts
    headers.pop("host", None)

    try:
        req = http_client.build_request(
            method=request.method,
            url=target_url,
            headers=headers,
            content=await request.body(),
            params=request.query_params,
        )
        resp = await http_client.send(req, stream=True)

        # Stream the response back
        return StreamingResponse(
            content=resp.aiter_raw(),
            status_code=resp.status_code,
            headers=dict(resp.headers),
        )

    except httpx.ConnectError:
        return StreamingResponse(
            content=iter([f"App '{app_name}' is not available".encode()]),
            status_code=503,
            media_type="text/plain",
        )
    except Exception:
        return StreamingResponse(
            content=iter([b"Internal proxy error"]),
            status_code=500,
            media_type="text/plain",
        )