#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Serve the bundled Asr-demo UI and bridge it to Qwen-ASR's realtime WS API.

Run from the Qwen-Asr repository root on the server where Qwen-ASR is listening:

    python scripts/test_qwen_ws_with_asr_demo.py \
      --backend-url ws://127.0.0.1:33050/ws/v1/asr/qwen

Then forward the UI port from your workstation with `ssh -L 8188:127.0.0.1:8188
user@server` and open http://127.0.0.1:8188. The bundled frontend means the
Asr-demo checkout is not needed on the server. The served HTML gets a small Qwen
label and file-type adjustment, while this bridge translates the legacy protocol.
Requires the `websockets` package, which is already listed by the project.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import mimetypes
import re
from http import HTTPStatus
from pathlib import Path
from typing import Any
from urllib.parse import unquote, urlsplit

try:
    from websockets.exceptions import ConnectionClosed
    from websockets.legacy.client import connect
    from websockets.legacy.server import serve
except ImportError as exc:
    raise SystemExit(
        "Missing websockets. Activate the Qwen-Asr or Asr-demo Python environment, "
        "or install it with `python -m pip install websockets`."
    ) from exc

LOGGER = logging.getLogger("qwen_ws_frontend_test")
DEFAULT_BACKEND_URL = "ws://127.0.0.1:33050/ws/v1/asr/qwen"
# Keep the test UI in-repository so a server deployment needs only this project.
DEFAULT_FRONTEND_DIR = Path(__file__).resolve().parent / "qwen_ws_test_frontend"


def _as_bool(value: Any) -> bool:
    """Normalize the 0/1 values emitted by the existing browser UI."""
    if isinstance(value, str):
        return value.strip().lower() not in {"", "0", "false", "no", "off"}
    return bool(value)


def _qwen_start_message(frontend_message: dict[str, Any]) -> dict[str, Any]:
    """Translate Asr-demo's flat start options into Qwen-ASR's payload shape."""
    source = str(frontend_message.get("source", "mic")).lower()
    filename = str(frontend_message.get("file_name", ""))
    suffix = Path(filename).suffix.lower()

    # Qwen's current websocket decoder accepts raw PCM16 or a WAV container.
    if source == "mic" or suffix == ".pcm":
        audio_format = "pcm"
    elif suffix == ".wav":
        audio_format = "wav"
    else:
        raise ValueError(
            "Qwen WebSocket test input must be microphone audio, .pcm, or .wav; "
            f"received {filename or source!r}."
        )

    payload: dict[str, Any] = {
        "format": audio_format,
        "sample_rate": 16000,
        "enable_speaker": _as_bool(frontend_message.get("speaker_diarization", 1)),
        "enable_speaker_identification": False,
        "enable_native_partial_stream": True,
    }

    # The reused UI exposes hotwords as a comma-separated string; Qwen accepts
    # this text through its realtime context prompt.
    hotwords = str(frontend_message.get("hotwords", "")).strip()
    if hotwords:
        payload["context"] = hotwords

    return {"type": "start", "payload": payload}


async def _relay_browser_to_qwen(browser: Any, qwen: Any) -> None:
    """Forward audio frames and adapt the legacy UI's control messages."""
    audio_format = "pcm"
    wav_buffer = bytearray()

    async for event in browser:
        if isinstance(event, str):
            message = json.loads(event)
            msg_type = message.get("type")
            if msg_type == "start":
                try:
                    message = _qwen_start_message(message)
                except ValueError as exc:
                    await browser.send(json.dumps({"type": "error", "message": str(exc)}))
                    return
                audio_format = message["payload"]["format"]
            elif msg_type == "eof":
                # The demo signals end-of-file; Qwen finalizes a stream on stop.
                if audio_format == "wav" and wav_buffer:
                    # Keep the WAV container intact; arbitrary network chunks do
                    # not each contain a WAV header and cannot be decoded alone.
                    await qwen.send(bytes(wav_buffer))
                    wav_buffer.clear()
                message = {"type": "stop"}
            elif msg_type == "stop" and audio_format == "wav" and wav_buffer:
                await qwen.send(bytes(wav_buffer))
                wav_buffer.clear()
            await qwen.send(json.dumps(message, ensure_ascii=False))
            continue

        if isinstance(event, bytes):
            if audio_format == "wav":
                wav_buffer.extend(event)
            else:
                await qwen.send(event)


async def _relay_qwen_to_browser(qwen: Any, browser: Any) -> None:
    """Pass Qwen-ASR's voice_id/start/sentences/end/error events to the UI."""
    async for message in qwen:
        await browser.send(message)


async def _bridge_websocket(browser: Any, backend_url: str) -> None:
    try:
        async with connect(backend_url, max_size=None) as qwen:
            LOGGER.info("Connected to Qwen-ASR WebSocket: %s", backend_url)
            tasks = {
                asyncio.create_task(_relay_browser_to_qwen(browser, qwen)),
                asyncio.create_task(_relay_qwen_to_browser(qwen, browser)),
            }
            done, pending = await asyncio.wait(
                tasks,
                return_when=asyncio.FIRST_COMPLETED,
            )
            for task in pending:
                task.cancel()
            await asyncio.gather(*pending, return_exceptions=True)

            # Surface real relay failures while treating disconnects as normal.
            for task in done:
                error = task.exception()
                if error and not isinstance(error, ConnectionClosed):
                    raise error
    except ConnectionClosed:
        pass
    except Exception as exc:
        LOGGER.exception("WebSocket bridge failed")
        try:
            await browser.send(json.dumps({"type": "error", "message": str(exc)}))
        except Exception:
            pass
    finally:
        try:
            await browser.close(code=1000)
        except Exception:
            pass


def _qwen_compatible_index(frontend_dir: Path) -> bytes:
    """Adjust labels and selectable formats in the served copy of index only."""
    html = (frontend_dir / "index.html").read_text(encoding="utf-8")
    html = html.replace(
        'id="engineModel" value="16k_zh_en_speaker" readonly',
        'id="engineModel" value="Qwen3-ASR realtime" readonly',
    )
    html = re.sub(
        r'(<input type="file" id="audioFile" accept=")[^"]*(" hidden>)',
        r"\1.pcm,.wav\2",
        html,
    )
    html = re.sub(
        r'(<span class="file-info" id="fileInfo">).*?(</span>)',
        r"\1Mic streams live; files: 16 kHz mono PCM or WAV (sent after upload).\2",
        html,
    )
    return html.encode("utf-8")


def _make_http_handler(frontend_dir: Path):
    async def process_request(path: str, headers: Any):
        parsed = urlsplit(path)
        if parsed.path == "/ws" and headers.get("Upgrade", "").lower() == "websocket":
            return None

        if parsed.path == "/ws":
            return (
                HTTPStatus.UPGRADE_REQUIRED,
                [("Content-Type", "text/plain")],
                b"WebSocket upgrade required",
            )

        if parsed.path == "/api/stop":
            # The UI also sends stop over WebSocket; avoid its legacy 404 fetch.
            return (
                HTTPStatus.OK,
                [("Content-Type", "application/json")],
                b'{"status":"ok","message":"Stop is sent over the WebSocket."}',
            )

        if parsed.path == "/":
            body = _qwen_compatible_index(frontend_dir)
            content_type = "text/html; charset=utf-8"
        else:
            asset_path = (frontend_dir / unquote(parsed.path).lstrip("/")).resolve()
            if not asset_path.is_relative_to(frontend_dir):
                return HTTPStatus.FORBIDDEN, [], b"Forbidden"
            if not asset_path.is_file():
                return HTTPStatus.NOT_FOUND, [], b"Not found"
            body = asset_path.read_bytes()
            content_type = mimetypes.guess_type(asset_path.name)[0] or "application/octet-stream"
            if content_type.startswith("text/") or content_type in {
                "application/javascript",
                "application/json",
            }:
                content_type += "; charset=utf-8"

        return HTTPStatus.OK, [("Content-Type", content_type)], body

    return process_request


async def run_server(frontend_dir: Path, backend_url: str, host: str, port: int) -> None:
    async def handle_frontend_socket(browser: Any) -> None:
        if urlsplit(browser.path).path != "/ws":
            await browser.close(code=1008, reason="Use /ws")
            return
        await _bridge_websocket(browser, backend_url)

    async with serve(
        handle_frontend_socket,
        host,
        port,
        process_request=_make_http_handler(frontend_dir),
        max_size=None,
    ):
        LOGGER.info("Open Asr-demo UI at http://127.0.0.1:%s", port)
        LOGGER.info("Forwarding realtime ASR to %s", backend_url)
        await asyncio.Future()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Serve the Asr-demo frontend and connect it to Qwen-ASR realtime WS."
    )
    parser.add_argument(
        "--backend-url",
        default=DEFAULT_BACKEND_URL,
        help=f"Qwen-ASR websocket URL (default: {DEFAULT_BACKEND_URL})",
    )
    parser.add_argument(
        "--frontend-dir",
        type=Path,
        default=DEFAULT_FRONTEND_DIR,
        help="Optional override for the bundled Asr-demo frontend directory.",
    )
    parser.add_argument(
        "--host",
        default="127.0.0.1",
        help="Local UI bind address. Keep localhost for browser microphone permission.",
    )
    parser.add_argument("--port", type=int, default=8188, help="Local UI HTTP port.")
    parser.add_argument(
        "--log-level",
        choices=("critical", "error", "warning", "info", "debug"),
        default="info",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    frontend_dir = args.frontend_dir.resolve()
    if not (frontend_dir / "index.html").is_file():
        raise SystemExit(
            f"Test frontend not found at {frontend_dir}. "
            "Restore scripts/qwen_ws_test_frontend or pass --frontend-dir."
        )

    logging.basicConfig(level=getattr(logging, args.log_level.upper()))
    try:
        asyncio.run(run_server(frontend_dir, args.backend_url, args.host, args.port))
    except KeyboardInterrupt:
        LOGGER.info("Frontend test bridge stopped.")


if __name__ == "__main__":
    main()