#!/usr/bin/env python3
import argparse
import asyncio
import json
from urllib.request import urlopen

from websockets.asyncio.client import connect


def fetch(base_url, path):
    url = f"{base_url.rstrip('/')}/{path.lstrip('/')}"
    with urlopen(url, timeout=10) as response:
        if response.status != 200:
            raise RuntimeError(f"{url} returned HTTP {response.status}")
        return response.read()


async def check_websocket(url):
    async with connect(url, open_timeout=10, close_timeout=5) as websocket:
        payload = json.loads(await asyncio.wait_for(websocket.recv(), timeout=10))
    expected = {"status": "ok", "channel_layer": "ok"}
    if payload != expected:
        raise RuntimeError(f"{url} returned unexpected payload: {payload}")


def main():
    parser = argparse.ArgumentParser(description="Check the deployed Hulumath Nginx surface")
    parser.add_argument("--base-url", required=True)
    parser.add_argument("--ws-url", required=True)
    parser.add_argument("--minimum-videos", type=int, default=54)
    args = parser.parse_args()

    health = json.loads(fetch(args.base_url, "/health/"))
    if (
        health.get("status") != "ok"
        or health.get("database") != "ok"
        or not health.get("version")
    ):
        raise RuntimeError(f"unexpected health payload: {health}")

    homepage = fetch(args.base_url, "/").decode("utf-8")
    if "葫芦数学" not in homepage:
        raise RuntimeError("homepage marker is missing")

    admin_login = fetch(args.base_url, "/admin/login/").decode("utf-8")
    if "进入运营后台" not in admin_login:
        raise RuntimeError("admin login marker is missing")

    catalog = json.loads(fetch(args.base_url, "/api/v1/content/videos/catalog/"))
    if catalog.get("total", 0) < args.minimum_videos:
        raise RuntimeError(f"video catalog is incomplete: {catalog.get('total', 0)}")

    asyncio.run(check_websocket(args.ws_url))
    print(
        "Production smoke checks passed: "
        f"health v{health['version']}, homepage, admin, "
        f"{catalog['total']} videos, websocket"
    )


if __name__ == "__main__":
    main()