#!/usr/bin/env python3
import json
import os
import urllib.parse
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

PROMETHEUS_URL = os.environ.get("PROMETHEUS_URL", "http://5.0.0.20:9090").rstrip("/")
SOURCE_URL = os.environ.get("SERVER_STATUS_SOURCE_URL", "").strip()
OUTPUT_PATH = Path(os.environ.get("SERVER_STATUS_JSON", "site/server-status.json"))
TIMEOUT_SECONDS = float(os.environ.get("PROMETHEUS_TIMEOUT_SECONDS", "5"))

def write_data(data):
    if not isinstance(data.get("servers"), list):
        raise SystemExit("server status JSON missing servers array")

    data["generated_at"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")

    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = OUTPUT_PATH.with_name(OUTPUT_PATH.name + ".tmp")
    tmp_path.write_text(json.dumps(data, separators=(",", ":")) + "\n")
    os.replace(tmp_path, OUTPUT_PATH)
    print(f"wrote {OUTPUT_PATH}")

def prom_value(query):
    url = PROMETHEUS_URL + "/api/v1/query?" + urllib.parse.urlencode({"query": query})
    with urllib.request.urlopen(url, timeout=TIMEOUT_SECONDS) as response:
        body = json.loads(response.read().decode("utf-8"))

    result = body.get("data", {}).get("result", [])
    if not result:
        return 0
    return int(float(result[0]["value"][1]))

def q(metric, labels):
    label_text = ",".join(f'{key}="{value}"' for key, value in labels.items())
    return f"sum({metric}{{{label_text}}}) or vector(0)"

def newserv(region, service, ship, version):
    return prom_value(q("pso_newserv_clients_connected", {
        "region": region,
        "service": service,
        "ship": ship,
        "version": version,
    }))

def adhoc(region, game):
    return prom_value(q("psppeeps_adhoc_connected_clients_by_product", {
        "region": region,
        "service": f"{region}-psppeeps-adhoc",
        "ship": "psp",
        "game": game,
    }))

def row(label, players):
    return {"label": label, "players": int(players)}

def main():
    if SOURCE_URL:
        with urllib.request.urlopen(SOURCE_URL, timeout=TIMEOUT_SECONDS) as response:
            write_data(json.loads(response.read().decode("utf-8")))
        return

    us_alis_v2 = newserv("us", "us-newserv-live", "live", "v2")
    us_alis_v3 = newserv("us", "us-newserv-live", "live", "v3")
    us_alis_bb = newserv("us", "us-newserv-live", "live", "v4")
    us_abion_hcbb = newserv("us", "us-newserv-hardcore", "hardcore", "v4")
    us_adhoc_psp1 = adhoc("us", "psp1")
    us_adhoc_psp2i = adhoc("us", "psp2i")

    eu_palma_v2 = newserv("eu", "eu-newserv-live", "live", "v2")
    eu_palma_v3 = newserv("eu", "eu-newserv-live", "live", "v3")
    eu_palma_bb = newserv("eu", "eu-newserv-live", "live", "v4")
    eu_aiedo_hcbb = newserv("eu", "eu-newserv-hardcore", "hardcore", "v4")
    eu_adhoc_psp1 = adhoc("eu", "psp1")
    eu_adhoc_psp2i = adhoc("eu", "psp2i")

    write_data({
        "servers": [
            {
                "label": "US Server",
                "players": us_alis_v2 + us_alis_v3 + us_alis_bb + us_abion_hcbb + us_adhoc_psp1 + us_adhoc_psp2i,
                "ships": [
                    {"label": "Alis", "rows": [row("V2", us_alis_v2), row("V3", us_alis_v3), row("BB", us_alis_bb)]},
                    {"label": "Abion", "rows": [row("HC/BB", us_abion_hcbb)]},
                    {"label": "AdHoc-US", "rows": [row("PSP1", us_adhoc_psp1), row("PSP2i", us_adhoc_psp2i)]},
                ],
            },
            {
                "label": "EU Server",
                "players": eu_palma_v2 + eu_palma_v3 + eu_palma_bb + eu_aiedo_hcbb + eu_adhoc_psp1 + eu_adhoc_psp2i,
                "ships": [
                    {"label": "Palma", "rows": [row("V2", eu_palma_v2), row("V3", eu_palma_v3), row("BB", eu_palma_bb)]},
                    {"label": "Aiedo", "rows": [row("HC/BB", eu_aiedo_hcbb)]},
                    {"label": "AdHoc-EU", "rows": [row("PSP1", eu_adhoc_psp1), row("PSP2i", eu_adhoc_psp2i)]},
                ],
            },
        ],
    })

if __name__ == "__main__":
    main()