#!/usr/bin/env python3
"""Materialize one verified Classic replay for the command-stream runtime.

The immutable serving index is the authority.  The source replay hash selects
an entry.  Replay and map bytes use content-addressed shared caches.  A small
session descriptor then lets the reusable engine launcher start the replay.
No replay-specific CRIU image is read or created.
"""

from __future__ import annotations

import argparse
import fcntl
import hashlib
import json
import os
import re
import shutil
import subprocess
import tempfile
import urllib.request
from pathlib import Path
from typing import Callable


SHA1 = re.compile(r"[0-9a-f]{40}\Z")
MAX_REPLAY_BYTES = 64 * 1024 * 1024
MAX_MAP_BYTES = 256 * 1024 * 1024
# Schema 3 (the R2 mirror, see deploy/r2-catalog-mirror.py) is purely
# content-addressed: object names ARE the sha1 of their bytes.
R2_REPLAY_KEY = re.compile(r"replays/corpus/[0-9a-f]{40}\.w3g\Z")
R2_MAP_KEY = re.compile(r"maps/[0-9a-f]{40}\.(w3m|w3x|tmp)\Z")
KNOWN_PROFILES = frozenset({
    "native-1100", "native-1110", "native-1120", "native-1140", "native-1220",
    "native-124ab", "native-124cde", "native-1285", "native-1311-d3d9"})


class CatalogMaterializeError(RuntimeError):
    pass


def sha1_file(path: Path, limit: int) -> tuple[str, int]:
    digest = hashlib.sha1()
    size = 0
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            size += len(chunk)
            if size > limit:
                raise CatalogMaterializeError("catalog object exceeds its limit")
            digest.update(chunk)
    if not size:
        raise CatalogMaterializeError("catalog object is empty")
    return digest.hexdigest(), size


def safe_object_key(value: object, *, map_object: bool = False) -> str:
    if not isinstance(value, str) or not value or len(value) > 512:
        raise CatalogMaterializeError("catalog object key is invalid")
    parts = value.split("/")
    if any(part in ("", ".", "..") for part in parts):
        raise CatalogMaterializeError("catalog object key escapes its namespace")
    allowed = value.startswith("maps/")
    if map_object:
        allowed = allowed or value.startswith(
            "worker-image/classic/replaykit/1.14/")
    if not allowed:
        raise CatalogMaterializeError("catalog object key is outside its namespace")
    return value


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 CatalogMaterializeError("catalog engine profile is unsupported")


def replay_object_key(source_sha1: str, played_sha1: str,
                      profile: str = "native-1285") -> str:
    if profile == "native-1140":
        return ("worker-image/classic/replaykit/1.14/v3/replays/" +
                f"{played_sha1}.w3g")
    suffix = "-c28" if played_sha1 != source_sha1 else ""
    return f"replays/archive/{source_sha1}{suffix}.w3g"


def default_fetcher(bucket: str, region: str) -> Callable[[str, Path], None]:
    def fetch(key: str, destination: Path) -> None:
        subprocess.run([
            "aws", "s3", "cp", f"s3://{bucket}/{key}", str(destination),
            "--only-show-errors", "--region", region,
        ], check=True, timeout=120)
    return fetch


def https_fetcher(base_url: str, limit: int = MAX_MAP_BYTES,
                  timeout: float = 60.0) -> Callable[[str, Path], None]:
    """Fetch content-addressed objects from the public CDN in front of R2.

    No credentials live on the CPU workers: the objects are public, immutable
    and named by their sha1, and ensure_cached() verifies that sha1 before a
    byte of them reaches a game directory.
    """
    base = base_url.rstrip("/")
    if not base.startswith("https://"):
        raise CatalogMaterializeError("catalog base URL must be https")

    def fetch(key: str, destination: Path) -> None:
        if not (R2_REPLAY_KEY.fullmatch(key) or R2_MAP_KEY.fullmatch(key)):
            raise CatalogMaterializeError("catalog object key is not content-addressed")
        request = urllib.request.Request(
            f"{base}/{key}", headers={"User-Agent": "w3cs-materialize/1"})
        size = 0
        with urllib.request.urlopen(request, timeout=timeout) as response, \
                destination.open("wb") as output:
            while chunk := response.read(1024 * 1024):
                size += len(chunk)
                if size > limit:
                    raise CatalogMaterializeError("catalog object exceeds its limit")
                output.write(chunk)
    return fetch


def ensure_cached(cache_path: Path, expected_sha1: str, limit: int,
                  key: str, fetch: Callable[[str, Path], None]) -> int:
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    lock_path = cache_path.with_suffix(cache_path.suffix + ".lock")
    with lock_path.open("a+b") as lock:
        fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
        if cache_path.is_file() and not cache_path.is_symlink():
            actual, size = sha1_file(cache_path, limit)
            if actual == expected_sha1:
                return size
            cache_path.unlink()
        descriptor, temporary_name = tempfile.mkstemp(
            prefix=f".{cache_path.name}.", suffix=".tmp",
            dir=cache_path.parent)
        os.close(descriptor)
        temporary = Path(temporary_name)
        try:
            fetch(key, temporary)
            actual, size = sha1_file(temporary, limit)
            if actual != expected_sha1:
                raise CatalogMaterializeError(
                    "downloaded catalog object has the wrong SHA-1")
            os.chmod(temporary, 0o444)
            os.replace(temporary, cache_path)
            return size
        finally:
            temporary.unlink(missing_ok=True)


def link_or_copy(source: Path, destination: Path) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_name(f".{destination.name}.{os.getpid()}.tmp")
    temporary.unlink(missing_ok=True)
    try:
        os.link(source, temporary)
    except OSError:
        shutil.copyfile(source, temporary)
    os.chmod(temporary, 0o444)
    os.replace(temporary, destination)


def materialize(index_path: Path, source_sha1: str, output_root: Path,
                fetch: Callable[[str, Path], None]) -> dict[str, object]:
    source_sha1 = source_sha1.lower()
    if not SHA1.fullmatch(source_sha1):
        raise CatalogMaterializeError("source replay SHA-1 is invalid")
    with index_path.open(encoding="utf-8") as source:
        index = json.load(source)
    schema = index.get("schema")
    if schema not in (2, 3) or not isinstance(index.get("replays"), dict):
        raise CatalogMaterializeError("Classic serving index is invalid")
    entry = index["replays"].get(source_sha1)
    if not isinstance(entry, dict):
        raise CatalogMaterializeError("replay is not in the verified catalog")
    if schema == 3:
        played_sha1 = str(entry.get("played") or "").lower()
        map_sha1 = str(entry.get("mapSha1") or "").lower()
        profile = str(entry.get("profile") or "")
        if profile not in KNOWN_PROFILES:
            raise CatalogMaterializeError("catalog engine profile is unsupported")
    else:
        played_sha1 = str(entry.get("playedReplaySha1") or "").lower()
        map_sha1 = str(entry.get("mapContentSha1") or "").lower()
        profile = engine_profile(entry)
    if not SHA1.fullmatch(played_sha1) or not SHA1.fullmatch(map_sha1):
        raise CatalogMaterializeError("catalog content identity is invalid")
    map_path = str(entry.get("mapPath") or "")
    map_parts = map_path.replace("/", "\\").split("\\")
    allow_expired = profile == "native-1140"
    suffixes = (".w3m", ".w3x", ".tmp") if allow_expired else (".w3m", ".w3x")
    if (len(map_parts) < 2 or map_parts[0].lower() != "maps" or
            any(part in ("", ".", "..") for part in map_parts) or
            not map_parts[-1].lower().endswith(suffixes)):
        raise CatalogMaterializeError("catalog map path is invalid")
    if schema == 3:
        map_suffix_hint = Path(map_parts[-1]).suffix.lower().lstrip(".")
        if str(entry.get("mapExt") or "") != map_suffix_hint:
            raise CatalogMaterializeError("catalog map extension is inconsistent")
        map_key = f"maps/{map_sha1}.{map_suffix_hint}"
        replay_key = f"replays/corpus/{played_sha1}.w3g"
    else:
        map_key = safe_object_key(entry.get("mapObjectKey"), map_object=True)
        replay_key = replay_object_key(source_sha1, played_sha1, profile)

    cache = output_root / ".content"
    replay_cache = cache / "replays" / f"{played_sha1}.w3g"
    map_suffix = Path(map_parts[-1]).suffix.lower()
    map_cache = cache / "maps" / f"{map_sha1}{map_suffix}"
    replay_bytes = ensure_cached(
        replay_cache, played_sha1, MAX_REPLAY_BYTES, replay_key, fetch)
    map_bytes = ensure_cached(
        map_cache, map_sha1, MAX_MAP_BYTES, map_key, fetch)

    replay_target = output_root / f"{source_sha1}.w3g"
    map_relative = Path("maps") / f"{map_sha1}{map_suffix}"
    map_target = output_root / map_relative
    link_or_copy(replay_cache, replay_target)
    link_or_copy(map_cache, map_target)
    metadata = {
        "schema": 1,
        "sourceReplaySha1": source_sha1,
        "replaySha1": played_sha1,
        "replayObjectKey": replay_key,
        "engineProfile": profile,
        "mapPath": "\\".join(map_parts),
        "mapFile": map_relative.as_posix(),
        "mapContentSha1": map_sha1,
        "mapObjectKey": map_key,
        "mapRequired": True,
        "allowExpiredPatchMap": allow_expired and map_suffix == ".tmp",
    }
    if schema == 3:
        mode = str(entry.get("mode") or "")
        if mode in ("roc", "tft"):
            metadata["patchFamily"] = mode
        players = [str(p) for p in (entry.get("players") or []) if str(p)]
        if players:
            metadata["displayName"] = " vs ".join(players[:2])
        if entry.get("durationMs"):
            metadata["durationMs"] = int(entry["durationMs"])
    metadata_path = output_root / f"{source_sha1}.json"
    descriptor, temporary_name = tempfile.mkstemp(
        prefix=f".{metadata_path.name}.", suffix=".tmp", dir=output_root)
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as output:
            json.dump(metadata, output, sort_keys=True, separators=(",", ":"))
            output.write("\n")
            output.flush()
            os.fsync(output.fileno())
        os.chmod(temporary_name, 0o444)
        os.replace(temporary_name, metadata_path)
    finally:
        try:
            os.unlink(temporary_name)
        except FileNotFoundError:
            pass
    return {**metadata, "replayBytes": replay_bytes, "mapBytes": map_bytes,
            "metadataPath": str(metadata_path),
            "replayPath": str(replay_target)}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--index", required=True, type=Path)
    parser.add_argument("--source-sha1", required=True)
    parser.add_argument("--output-root", required=True, type=Path)
    parser.add_argument("--bucket")
    parser.add_argument("--region", default="us-east-1")
    parser.add_argument("--base-url",
                        help="public CDN base of the R2 mirror (schema 3 index)")
    args = parser.parse_args()
    if bool(args.bucket) == bool(args.base_url):
        parser.error("give exactly one of --bucket or --base-url")
    args.output_root.mkdir(parents=True, exist_ok=True)
    fetch = (https_fetcher(args.base_url) if args.base_url
             else default_fetcher(args.bucket, args.region))
    result = materialize(args.index, args.source_sha1, args.output_root, fetch)
    print(json.dumps(result, sort_keys=True))


if __name__ == "__main__":
    main()