#!/usr/bin/env python3
"""Serve the unchanged Tencent demo UI and proxy its API to the backend."""

from __future__ import annotations

import asyncio
import os
from pathlib import Path
from urllib.parse import urlsplit
from typing import Any

from aiohttp import ClientSession, ClientTimeout, WSMsgType, web
from dotenv import load_dotenv

PROJECT_ROOT = Path(__file__).resolve().parents[1]
STATIC_ROOT = Path(__file__).resolve().parent / "static"
load_dotenv(PROJECT_ROOT / ".env")
FRONTEND_HOST = os.getenv("FRONTEND_HOST", "127.0.0.1")
FRONTEND_PORT = int(os.getenv("FRONTEND_PORT", "8080"))
BACKEND_BASE_URL = os.getenv(
    "BACKEND_INTERNAL_URL",
    f"http://127.0.0.1:{os.getenv('WEB_PORT', '8082')}",
).rstrip("/")
HTTP_SESSION = web.AppKey("http_session", ClientSession)


def backend_url(request: web.Request) -> str:
    """Keep the original path and query while routing through the backend port."""
    return f"{BACKEND_BASE_URL}{request.rel_url}"


async def index_handler(_: web.Request) -> web.FileResponse:
    return web.FileResponse(
        STATIC_ROOT / "index.html", headers={"Cache-Control": "no-store"}
    )


async def api_stop_proxy(request: web.Request) -> web.Response:
    """Forward the Tencent page's existing stop request to the WS backend."""
    async with request.app[HTTP_SESSION].get(
        backend_url(request), timeout=ClientTimeout(total=5)
    ) as response:
        body = await response.read()
        return web.Response(
            status=response.status,
            body=body,
            headers={"Content-Type": response.headers.get("Content-Type", "application/json")},
        )


async def websocket_proxy(request: web.Request) -> web.WebSocketResponse:
    """Relay text and binary frames without changing the Tencent browser protocol."""
    browser_ws = web.WebSocketResponse(max_msg_size=64 * 1024 * 1024, heartbeat=30)
    await browser_ws.prepare(request)
    try:
        backend_ws = await request.app[HTTP_SESSION].ws_connect(
            backend_url(request),
            max_msg_size=64 * 1024 * 1024,
            heartbeat=30,
            autoping=True,
        )
    except Exception as exc:
        await browser_ws.send_json({"type": "error", "message": f"backend unavailable: {exc}"})
        await browser_ws.close(code=1011, message=b"backend unavailable")
        return browser_ws

    async def relay(source: Any, destination: Any) -> None:
        async for message in source:
            if message.type == WSMsgType.TEXT:
                await destination.send_str(message.data)
            elif message.type == WSMsgType.BINARY:
                await destination.send_bytes(message.data)
            elif message.type in {WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.ERROR}:
                if not destination.closed:
                    code = message.data if isinstance(message.data, int) else 1000
                    reason = message.extra or ""
                    await destination.close(code=code, message=str(reason).encode("utf-8"))
                return

    tasks = [
        asyncio.create_task(relay(browser_ws, backend_ws)),
        asyncio.create_task(relay(backend_ws, browser_ws)),
    ]
    try:
        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        for task in pending:
            task.cancel()
        await asyncio.gather(*done, *pending, return_exceptions=True)
    finally:
        if not backend_ws.closed:
            await backend_ws.close()
        if not browser_ws.closed:
            await browser_ws.close()
    return browser_ws


async def create_app() -> web.Application:
    if not STATIC_ROOT.is_dir():
        raise FileNotFoundError(f"Tencent demo static directory is missing: {STATIC_ROOT}")
    parsed = urlsplit(BACKEND_BASE_URL)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc or parsed.path:
        raise ValueError("BACKEND_INTERNAL_URL must contain only an HTTP origin")

    app = web.Application()

    async def lifecycle(application: web.Application):
        # An unbounded total timeout allows long recordings and slow model loads.
        application[HTTP_SESSION] = ClientSession(
            timeout=ClientTimeout(total=None, connect=10, sock_connect=10, sock_read=None)
        )
        yield
        await application[HTTP_SESSION].close()

    app.cleanup_ctx.append(lifecycle)
    app.router.add_get("/", index_handler)
    app.router.add_get("/ws", websocket_proxy)
    app.router.add_get("/api/stop", api_stop_proxy)
    app.router.add_static("/", STATIC_ROOT, show_index=False)
    return app


def main() -> None:
    print(
        f"Tencent demo frontend: http://{FRONTEND_HOST}:{FRONTEND_PORT}/ "
        f"(backend proxy: {BACKEND_BASE_URL})",
        flush=True,
    )
    web.run_app(create_app(), host=FRONTEND_HOST, port=FRONTEND_PORT)


if __name__ == "__main__":
    main()