#!/usr/bin/env python3
"""Plan and publish the R2 catalog mirror for the CPU workers.

The CPU workers (London) fetch replays and maps on demand from R2 through
the public CDN host, never from AWS. R2 is content-addressed:

    replays/corpus/<sha1-of-bytes>.w3g   originals AND c28 conversions
    maps/<sha1-of-bytes>.<w3m|w3x|tmp>   every map the catalog can reference
    catalog/current.json                 -> {"sha256","key","replays",...}
    catalog/serving-index-<sha256>.json  compact immutable serving index

Subcommands (run from the developer machine; AWS creds from the
mastering-aws profile, R2 creds from the repo .env):

  plan     write a prioritized plan.jsonl of S3 objects still missing from
           R2, each with a time-limited presigned GET URL, so a box with no
           AWS credentials can do the transfer (r2-mirror-worker.py)
  index    compact the verified serving index (schema 2, the AWS layer
           index) into the schema-3 index the workers read
  publish  upload the compact index to R2 and flip catalog/current.json
"""
from __future__ import annotations

import argparse
import gzip
import hashlib
import json
import os
import sys
import time
from pathlib import Path

SHA1_HEX = 40


def r2_client():
    import boto3
    env = {}
    env_path = Path(__file__).resolve().parents[3] / ".env"
    for line in env_path.read_text().splitlines():
        if "=" in line and not line.startswith("#"):
            key, _, value = line.partition("=")
            env[key.strip()] = value.strip().strip('"').strip("'")
    session = boto3.session.Session(
        aws_access_key_id=env["R2_ACCESS_KEY_ID"],
        aws_secret_access_key=env["R2_SECRET_ACCESS_KEY"],
        region_name="auto")
    # boto3 >= 1.36 adds CRC32 checksum trailers by default; R2 rejects or
    # stalls on them. Only compute checksums where the API requires them.
    from botocore.config import Config
    config = Config(request_checksum_calculation="when_required",
                    response_checksum_validation="when_required",
                    retries={"max_attempts": 3})
    return (session.client("s3", endpoint_url=env["R2_ENDPOINT"], config=config),
            env["R2_BUCKET"])


def s3_client(profile: str, region: str):
    import boto3
    return boto3.session.Session(profile_name=profile, region_name=region).client("s3")


def list_keys(client, bucket: str, prefix: str) -> set[str]:
    keys: set[str] = set()
    paginator = client.get_paginator("list_objects_v2")
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
        for obj in page.get("Contents", []):
            keys.add(obj["Key"])
    return keys


def map_extension(name: str) -> str:
    ext = name.rsplit(".", 1)[-1].lower() if "." in name else ""
    if ext not in ("w3m", "w3x", "tmp"):
        raise ValueError(f"unsupported map extension in {name!r}")
    return ext


def engine_profile(entry: dict) -> str:
    base = str(entry.get("baseId") or "")
    engine = str(entry.get("engine") or "")
    if base.startswith("classic-1140-") or engine == "1.14":
        return "native-1140"
    if base.startswith("classic-1285-") or engine == "1.28.5":
        return "native-1285"
    if base.startswith("classic-1311-") or engine == "1.31.1":
        return "native-1311-d3d9"
    raise ValueError(f"unsupported engine {engine!r} / {base!r}")


def compact_index(layer_index: dict) -> dict:
    """Schema 3: only what a worker needs to materialize a replay from R2."""
    if layer_index.get("schema") != 2:
        raise ValueError("expected the schema-2 layer index")
    out = {}
    for source, entry in sorted(layer_index["replays"].items()):
        played = str(entry["playedReplaySha1"]).lower()
        map_sha1 = str(entry["mapContentSha1"]).lower()
        if len(source) != SHA1_HEX or len(played) != SHA1_HEX or len(map_sha1) != SHA1_HEX:
            raise ValueError(f"bad identity in {source}")
        map_path = str(entry["mapPath"])
        mode = str(entry.get("classicGameMode") or "").lower()
        if mode not in ("roc", "tft"):
            raise ValueError(f"bad game mode in {source}")
        out[source] = {
            "played": played,
            "mapSha1": map_sha1,
            "mapExt": map_extension(map_path.replace("/", "\\").split("\\")[-1]),
            "mapPath": map_path,
            "profile": engine_profile(entry),
            "mode": mode,
            "family": str(entry.get("family") or mode.upper()),
            "players": [str(p) for p in (entry.get("players") or [])][:12],
            "durationMs": int(entry.get("durationMs") or 0),
        }
    return {"schema": 3, "generatedAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "replayCount": len(out), "replays": out}


def cmd_index(args):
    layer = json.loads(Path(args.layer_index).read_text())
    compact = compact_index(layer)
    body = json.dumps(compact, sort_keys=True, separators=(",", ":")).encode()
    Path(args.output).write_bytes(body)
    print(json.dumps({"replays": compact["replayCount"], "bytes": len(body),
                      "sha256": hashlib.sha256(body).hexdigest(),
                      "output": args.output}))


def cmd_publish(args):
    body = Path(args.index).read_bytes()
    data = json.loads(body)
    if data.get("schema") != 3:
        raise SystemExit("publish expects a schema-3 index")
    sha256 = hashlib.sha256(body).hexdigest()
    client, bucket = r2_client()
    key = f"catalog/serving-index-{sha256}.json"
    client.put_object(Bucket=bucket, Key=key, Body=body,
                      ContentType="application/json",
                      CacheControl="public, max-age=31536000, immutable")
    current = {"schema": 1, "sha256": sha256, "key": key,
               "replays": data["replayCount"], "generatedAt": data["generatedAt"],
               "publishedAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}
    client.put_object(Bucket=bucket, Key="catalog/current.json",
                      Body=json.dumps(current, sort_keys=True).encode(),
                      ContentType="application/json",
                      CacheControl="public, max-age=60")
    print(json.dumps(current))


def cmd_plan(args):
    s3 = s3_client(args.profile, args.region)
    r2, r2_bucket = r2_client()
    layer = json.loads(Path(args.layer_index).read_text())["replays"]
    manifest = json.loads(Path(args.maps_manifest).read_text())
    readiness = [json.loads(line) for line in gzip.open(args.readiness, "rt")]

    print("listing R2 ...", file=sys.stderr)
    r2_replays = {k.split("/")[-1][:-4] for k in list_keys(r2, r2_bucket, "replays/corpus/")
                  if k.endswith(".w3g")}
    r2_maps = {k.split("/")[-1].split(".")[0] for k in list_keys(r2, r2_bucket, "maps/")}
    print(f"R2 has {len(r2_replays)} replays, {len(r2_maps)} maps", file=sys.stderr)

    def presign(key: str) -> str:
        return s3.generate_presigned_url(
            "get_object", Params={"Bucket": args.bucket, "Key": key},
            ExpiresIn=args.expires)

    plan: list[dict] = []
    seen: set[str] = set()

    def add(kind: str, s3_key: str, r2_key: str | None, sha1: str | None, priority: int):
        ident = r2_key or f"auto:{s3_key}"
        if ident in seen:
            return
        seen.add(ident)
        plan.append({"priority": priority, "kind": kind, "s3Key": s3_key,
                     "key": r2_key, "sha1": sha1, "url": presign(s3_key)})

    # 1. maps the verified index references, 2. their conversions.
    index_map_keys: dict[str, str] = {}
    for source, entry in layer.items():
        index_map_keys.setdefault(entry["mapContentSha1"].lower(), entry["mapObjectKey"])
    for map_sha1, s3_key in index_map_keys.items():
        if map_sha1 in r2_maps:
            continue
        add("map", s3_key, f"maps/{map_sha1}.{map_extension(s3_key)}", map_sha1, 1)
    for source, entry in layer.items():
        played = entry["playedReplaySha1"].lower()
        if played in r2_replays:
            continue
        s3_key = (f"replays/archive/{source}-c28.w3g" if played != source
                  else f"replays/archive/{source}.w3g")
        if entry.get("engine") == "1.14":
            s3_key = f"worker-image/classic/replaykit/1.14/v3/replays/{played}.w3g"
        add("replay", s3_key, f"replays/corpus/{played}.w3g", played, 2)
    # 3. every other map content (the full corpus, so the next release needs
    # no map transfer), 4. every other conversion and original.
    for map_sha1, keys in manifest["content"].items():
        if map_sha1 in r2_maps or map_sha1 in index_map_keys:
            continue
        s3_key = sorted(keys)[0]
        try:
            ext = map_extension(s3_key)
        except ValueError:
            continue
        add("map", s3_key, f"maps/{map_sha1}.{ext}", map_sha1, 3)
    for entry in readiness:
        source = entry["sha1"].lower()
        if entry.get("convertedPresent") and entry.get("convertedKey") and source not in layer:
            add("replay", entry["convertedKey"], None, None, 4)   # sha1 known after download
        if entry.get("originalPresent") and source not in r2_replays:
            add("replay", entry["originalKey"], f"replays/corpus/{source}.w3g", source, 5)

    plan.sort(key=lambda item: item["priority"])
    with open(args.output, "w") as out:
        for item in plan:
            out.write(json.dumps(item) + "\n")
    counts = {}
    for item in plan:
        counts[item["priority"]] = counts.get(item["priority"], 0) + 1
    print(json.dumps({"planned": len(plan), "byPriority": counts,
                      "expiresInHours": args.expires / 3600, "output": args.output}))


def main():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="command", required=True)
    p = sub.add_parser("plan")
    p.add_argument("--layer-index", required=True)
    p.add_argument("--maps-manifest", required=True)
    p.add_argument("--readiness", required=True)
    p.add_argument("--output", required=True)
    p.add_argument("--bucket", default="war3replays-assets-440964995465")
    p.add_argument("--profile", default="mastering-aws")
    p.add_argument("--region", default="us-east-1")
    p.add_argument("--expires", type=int, default=48 * 3600)
    p.set_defaults(func=cmd_plan)
    p = sub.add_parser("index")
    p.add_argument("--layer-index", required=True)
    p.add_argument("--output", required=True)
    p.set_defaults(func=cmd_index)
    p = sub.add_parser("publish")
    p.add_argument("--index", required=True)
    p.set_defaults(func=cmd_publish)
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()