#!/usr/bin/env python3
"""Reconstruct the real C3S seasonal release calendar from the CDS object store.

The CDS regenerates a dataset's `constraints.json` every time new data is
registered, and writes it to a public, content-addressed object store that
keeps every past version with its write timestamp. Listing that bucket
therefore recovers the exact moment each batch of forecasts became available --
no API key, no waiting, no polling.

    python release_schedule.py                    # calendar + latency stats
    python release_schedule.py --diff             # also: what each batch added
    python release_schedule.py --dataset seasonal-monthly-single-levels

Latency is measured against the 1st of the month at 00 UTC, which is the
initialisation time every system in the multi-system shares.
"""

import argparse
import datetime as dt
import json
import statistics
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

BUCKET = "https://object-store.os-api.cci2.ecmwf.int/cci2-prod-catalogue"
NS = "{http://s3.amazonaws.com/doc/2006-03-01/}"
VARIABLE = "maximum_2m_temperature_in_the_last_24_hours"


def list_versions(dataset):
    """Every stored constraints.json for this dataset, oldest first."""
    prefix = f"resources/{dataset}/constraints"
    out, marker = [], ""
    while True:
        q = urllib.parse.urlencode({"prefix": prefix, "max-keys": "1000", "marker": marker})
        with urllib.request.urlopen(f"{BUCKET}?{q}", timeout=120) as fh:
            root = ET.fromstring(fh.read())
        got = 0
        for c in root.findall(NS + "Contents"):
            key = c.find(NS + "Key").text
            stamp = dt.datetime.strptime(
                c.find(NS + "LastModified").text, "%Y-%m-%dT%H:%M:%S.%fZ"
            ).replace(tzinfo=dt.timezone.utc)
            out.append((stamp, key))
            marker = key
            got += 1
        truncated = root.find(NS + "IsTruncated")
        if got == 0 or truncated is None or truncated.text != "true":
            break
    return sorted(out)


def fetch(key):
    with urllib.request.urlopen(f"{BUCKET}/{urllib.parse.quote(key)}", timeout=300) as fh:
        return json.load(fh)


def offered(constraints):
    """{(centre, system): {(year, month, day)}} for the EHD variable."""
    out = {}
    for b in constraints:
        if VARIABLE not in b.get("variable", ()):
            continue
        dates = {(y, m, d) for y in b["year"] for m in b["month"] for d in b["day"]}
        for centre in b["originating_centre"]:
            for system in b["system"]:
                out.setdefault((centre, system), set()).update(dates)
    return out


def latency_days(stamp):
    """Gap between the 1st of the release month at 00 UTC and the write time."""
    init = dt.datetime(stamp.year, stamp.month, 1, tzinfo=dt.timezone.utc)
    return (stamp - init).total_seconds() / 86400


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--dataset", default="seasonal-original-single-levels")
    ap.add_argument("--since", default="2025-01-01", help="ISO date; earlier writes are noisy")
    ap.add_argument("--diff", action="store_true",
                    help="download consecutive versions and report what each batch added")
    args = ap.parse_args()

    since = dt.datetime.fromisoformat(args.since).replace(tzinfo=dt.timezone.utc)
    versions = list_versions(args.dataset)
    recent = [(s, k) for s, k in versions if s >= since]

    print(f"{args.dataset}: {len(versions)} phien ban luu tru, "
          f"{len(recent)} tu {args.since}\n")

    print(f"{'ghi luc (UTC)':<22}{'ngay':>6}{'do tre so voi init ngay 01':>30}")
    for stamp, _ in recent:
        print(f"{stamp:%Y-%m-%d %H:%M:%S}  {stamp.day:>4}"
              f"{latency_days(stamp):>26.3f} ngay")

    groups = {}
    for stamp, _ in recent:
        groups.setdefault(stamp.day, []).append(stamp)
    print("\nThong ke theo ngay trong thang:")
    for day in sorted(groups):
        stamps = groups[day]
        lat = [latency_days(s) for s in stamps]
        tod = [s.hour * 60 + s.minute for s in stamps]
        tag = "lich dinh ky" if len(stamps) > 2 else "ngoai lich"
        print(f"  ngay {day:>2} ({tag}, n={len(stamps)}): "
              f"do tre {min(lat):.3f} - {max(lat):.3f} ngay, "
              f"gio {min(tod)//60:02d}:{min(tod)%60:02d} - {max(tod)//60:02d}:{max(tod)%60:02d}")

    if not args.diff:
        print("\n(them --diff de xem moi lo mang them nhung gi -- se tai vai chuc MB)")
        return 0

    print("\nNoi dung moi cua tung lo:")
    prev = None
    for stamp, key in recent:
        cur = offered(fetch(key))
        if prev is not None:
            added = {}
            for sysid, dates in cur.items():
                new = dates - prev.get(sysid, set())
                if new:
                    added[sysid] = new
            if added:
                print(f"\n  {stamp:%Y-%m-%d %H:%M UTC}")
                for sysid in sorted(added):
                    months = {}
                    for y, m, d in added[sysid]:
                        months.setdefault(f"{y}-{m}", 0)
                        months[f"{y}-{m}"] += 1
                    detail = ", ".join(f"{k}: {v} ngay" for k, v in sorted(months.items()))
                    print(f"    {sysid[0]}/{sysid[1]:<5} {detail}")
        prev = cur
    return 0


if __name__ == "__main__":
    raise SystemExit(main())