#!/usr/bin/env python3
"""Execute an R2 mirror plan on a box that has rclone + R2 write access.

Each plan line names one S3 object (as a presigned GET URL), its expected
sha1 (or null: content-addressed after download) and its R2 key. The worker
downloads with a thread pool, verifies the sha1, stages the file under its
final key name and pushes batches with one `rclone copy` each. Immutable
objects get a one-year cache header. Progress and failures are append-only
logs, so a rerun resumes where it stopped.

    sudo r2-mirror-worker.py --plan plan.jsonl --staging /var/tmp/w3cs-r2
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import subprocess
import sys
import threading
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

MAX_BYTES = 256 * 1024 * 1024
CACHE = "Cache-Control: public, max-age=31536000, immutable"


def download(item: dict, staging: Path, retries: int = 3) -> dict:
    for attempt in range(retries):
        try:
            temp = staging / ".tmp" / f"{threading.get_ident()}-{os.getpid()}-{attempt}.part"
            temp.parent.mkdir(parents=True, exist_ok=True)
            digest = hashlib.sha1()
            size = 0
            with urllib.request.urlopen(item["url"], timeout=60) as response, temp.open("wb") as out:
                while chunk := response.read(1024 * 1024):
                    size += len(chunk)
                    if size > MAX_BYTES:
                        raise ValueError("object exceeds the size limit")
                    digest.update(chunk)
                    out.write(chunk)
            sha1 = digest.hexdigest()
            if not size:
                raise ValueError("empty object")
            if item.get("sha1") and sha1 != item["sha1"]:
                raise ValueError(f"sha1 mismatch: {sha1} != {item['sha1']}")
            key = item.get("key") or f"replays/corpus/{sha1}.w3g"
            final = staging / key
            final.parent.mkdir(parents=True, exist_ok=True)
            os.replace(temp, final)
            return {"key": key, "sha1": sha1, "bytes": size, "s3Key": item["s3Key"]}
        except Exception as error:  # noqa: BLE001
            last = error
            time.sleep(1 + attempt)
    return {"error": str(last), "s3Key": item["s3Key"], "key": item.get("key")}


def push_batch(staging: Path, remote: str, config: str, transfers: int) -> None:
    for prefix in ("replays/corpus", "maps"):
        source = staging / prefix
        if not source.is_dir() or not any(source.iterdir()):
            continue
        subprocess.run([
            "rclone", "--config", config, "copy", str(source), f"{remote}/{prefix}",
            "--ignore-existing", "--no-traverse", "--transfers", str(transfers),
            "--checkers", str(transfers), "--retries", "3",
            "--header-upload", CACHE, "--stats-one-line", "--stats", "0",
        ], check=True, timeout=3600)
        shutil.rmtree(source)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--plan", required=True)
    parser.add_argument("--staging", default="/var/tmp/w3cs-r2-mirror")
    parser.add_argument("--rclone-config", default="/etc/w3cs/rclone.conf")
    parser.add_argument("--remote", default="r2:war3replays")
    parser.add_argument("--threads", type=int, default=24)
    parser.add_argument("--transfers", type=int, default=16)
    parser.add_argument("--batch", type=int, default=1500)
    args = parser.parse_args()

    staging = Path(args.staging)
    staging.mkdir(parents=True, exist_ok=True)
    done_path = staging / "done.jsonl"
    failed_path = staging / "failed.jsonl"
    done_ids: set[str] = set()
    if done_path.exists():
        for line in done_path.read_text().splitlines():
            if line.strip():
                done_ids.add(json.loads(line)["s3Key"])
    items = [json.loads(line) for line in open(args.plan) if line.strip()]
    pending = [item for item in items if item["s3Key"] not in done_ids]
    print(f"plan {len(items)} items, {len(done_ids)} done, {len(pending)} pending", flush=True)
    started = time.time()
    total_bytes = 0
    completed = 0
    failures = 0
    with ThreadPoolExecutor(max_workers=args.threads) as pool:
        for start in range(0, len(pending), args.batch):
            batch = pending[start:start + args.batch]
            results = [future.result() for future in
                       as_completed([pool.submit(download, item, staging) for item in batch])]
            ok = [r for r in results if "error" not in r]
            bad = [r for r in results if "error" in r]
            push_batch(staging, args.remote, args.rclone_config, args.transfers)
            with done_path.open("a") as out:
                for r in ok:
                    out.write(json.dumps(r) + "\n")
            with failed_path.open("a") as out:
                for r in bad:
                    out.write(json.dumps(r) + "\n")
            completed += len(ok)
            failures += len(bad)
            total_bytes += sum(r["bytes"] for r in ok)
            elapsed = time.time() - started
            print(f"{completed + failures}/{len(pending)} "
                  f"ok={completed} failed={failures} "
                  f"{total_bytes / 1e9:.2f} GB {elapsed / 60:.1f} min "
                  f"{total_bytes / max(elapsed, 1) / 1e6:.1f} MB/s", flush=True)
    print("MIRROR-DONE", json.dumps({"ok": completed, "failed": failures,
                                     "bytes": total_bytes}), flush=True)


if __name__ == "__main__":
    main()