"""Receive China airport METAR updates and recover latest state after reconnecting.

Python 3.11+; install with: python -m pip install httpx "websockets>=15"
Set SKYVIEWOR_API_KEY in your shell. Never put the key in this file.
"""

import argparse
import asyncio
import json
import os
import random
import re
from urllib.parse import urlencode

import httpx
from websockets.asyncio.client import connect
from websockets.exceptions import ConnectionClosed, InvalidStatus

WS_URL = "wss://data.skyviewor.io/v1/ws/metar"
LATEST_URL = "https://data.skyviewor.io/v1/metar/latest"
MAX_ATTEMPTS = 8


def airport_list(value: str) -> str:
    airports = [item.strip().upper() for item in value.split(",")]
    if not airports or any(not re.fullmatch(r"[A-Z]{4}", item) for item in airports):
        raise argparse.ArgumentTypeError("Use comma-separated four-letter ICAO codes, such as ZBAA,ZSPD")
    return ",".join(airports)


def latest(airports: str, key: str) -> dict:
    with httpx.Client(timeout=15.0) as client:
        response = client.get(
            LATEST_URL,
            params={"airports": airports},
            headers={"Authorization": f"Bearer {key}"},
        )
        if response.status_code == 404:
            return {"data": []}
        response.raise_for_status()
        return response.json()


async def stream_once(airports: str, key: str) -> None:
    url = f"{WS_URL}?{urlencode({'airports': airports})}"
    async with connect(
        url,
        additional_headers={"Authorization": f"Bearer {key}"},
        proxy=None,
        open_timeout=15,
        ping_interval=20,
        ping_timeout=20,
        close_timeout=5,
    ) as websocket:
        print("Connected. Fetching latest state to cover the reconnect gap.", flush=True)
        snapshot = await asyncio.to_thread(latest, airports, key)
        print("LATEST " + json.dumps(snapshot, ensure_ascii=False), flush=True)
        pong = asyncio.Event()

        async def heartbeat() -> None:
            while True:
                started = asyncio.get_running_loop().time()
                pong.clear()
                await asyncio.wait_for(websocket.send('{"action":"ping"}'), 10)
                await asyncio.wait_for(pong.wait(), 15)
                print("Heartbeat OK", flush=True)
                await asyncio.sleep(max(0, 20 - (asyncio.get_running_loop().time() - started)))

        async def receive() -> None:
            async for raw in websocket:
                message = json.loads(raw)
                if message.get("type") == "pong":
                    pong.set()
                elif message.get("type") == "new_metar_arrived":
                    print("UPDATE " + json.dumps(message.get("data"), ensure_ascii=False), flush=True)
                elif message.get("type") == "error":
                    raise RuntimeError("WebSocket reported a subscription error; see the METAR API reference")

        tasks = [asyncio.create_task(receive()), asyncio.create_task(heartbeat())]
        try:
            done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for task in done:
                task.result()
        finally:
            for task in tasks:
                task.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)
    # The old connection is closed before another connection is attempted.


async def run(airports: str, key: str, seconds: int) -> None:
    attempts = 0
    try:
        async with asyncio.timeout(seconds):
            while attempts < MAX_ATTEMPTS:
                try:
                    await stream_once(airports, key)
                except InvalidStatus as exc:
                    if exc.response.status_code in (401, 403):
                        raise SystemExit("Authentication failed; check your key and trial status") from None
                    if exc.response.status_code != 429:
                        raise SystemExit(f"WebSocket handshake returned HTTP {exc.response.status_code}") from None
                    print("Connection quota is occupied (HTTP 429).", flush=True)
                except httpx.HTTPStatusError as exc:
                    if exc.response.status_code in (401, 403):
                        raise SystemExit("Latest-state authentication failed; check your key and trial status") from None
                    print(f"Latest-state request returned HTTP {exc.response.status_code}", flush=True)
                except (ConnectionClosed, OSError, httpx.HTTPError, asyncio.TimeoutError) as exc:
                    print(f"Connection interrupted: {type(exc).__name__}", flush=True)
                attempts += 1
                if attempts == MAX_ATTEMPTS:
                    break
                delay = min(20, 2 ** (attempts - 1)) + random.uniform(0, 0.5)
                print(f"Retrying serially in {delay:.1f}s", flush=True)
                await asyncio.sleep(delay)
    except TimeoutError:
        print("Demo time ended. Heartbeats prove connection health; an update requires a new report.")
        return
    raise SystemExit("Reconnect budget exhausted. Check other connections sharing this key and try later.")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--airports", type=airport_list, default="ZBAA,ZSPD")
    parser.add_argument("--seconds", type=int, default=180, help="Total demo duration (default: 180)")
    args = parser.parse_args()
    if not 1 <= args.seconds <= 3600:
        parser.error("--seconds must be between 1 and 3600")
    key = os.getenv("SKYVIEWOR_API_KEY")
    if not key:
        parser.error("Set SKYVIEWOR_API_KEY in your environment before making a live request")
    try:
        asyncio.run(run(args.airports, key, args.seconds))
    except KeyboardInterrupt:
        print("Stopped")


if __name__ == "__main__":
    main()