"""Archive a licensed, public unbundled SOG; no cookies or account tokens.

Use only a metadata URL observed in an authorized scene. This downloads assets,
not collision/voxel geometry. The source provenance is carried into the manifest.
"""

import argparse
import gzip
import hashlib
import json
import re
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from datetime import UTC, datetime
from pathlib import Path
from urllib.parse import urljoin, urlsplit

parser = argparse.ArgumentParser()
parser.add_argument("--metadata-url", required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--author", required=True)
parser.add_argument("--license", required=True)
parser.add_argument("--source-url", required=True)
args = parser.parse_args()
url = urlsplit(args.metadata_url)
if url.scheme != "https" or url.username or not url.path.endswith("/meta.json"):
    parser.error("Use a public HTTPS SOG meta.json URL")
args.output.mkdir(parents=True, exist_ok=False)
started = time.monotonic_ns()


def fetch(name, limit):
    target = args.output / name
    digest = hashlib.sha256()
    length = 0
    with (
        urllib.request.urlopen(urljoin(args.metadata_url, name), timeout=30) as response,
        target.with_suffix(target.suffix + ".part").open("xb") as stream,
    ):
        while data := response.read(1024 * 1024):
            length += len(data)
            if length > limit:
                raise ValueError("SOG asset exceeds bounded download size")
            stream.write(data)
            digest.update(data)
    target.with_suffix(target.suffix + ".part").replace(target)
    return {"path": name, "byte_length": length, "sha256": digest.hexdigest()}


meta = fetch("meta.json", 1024 * 1024)
raw = (args.output / "meta.json").read_bytes()
if raw[:2] == b"\x1f\x8b":
    raw = gzip.decompress(raw)
    if len(raw) > 1024 * 1024:
        raise ValueError("SOG metadata exceeds bounded size")
    (args.output / "meta.json").write_bytes(raw)
    meta = {"path": "meta.json", "byte_length": len(raw), "sha256": hashlib.sha256(raw).hexdigest()}
metadata = json.loads(raw)
if metadata.get("version") != 2 or not 1 <= metadata["count"] <= 20_000_000:
    raise ValueError("Unsupported SOG version or size")
files = sorted(
    {
        name
        for value in metadata.values()
        if isinstance(value, dict)
        for name in value.get("files", [])
    }
)
if not 1 <= len(files) <= 10 or any(not re.fullmatch(r"[A-Za-z0-9_-]+\.webp", n) for n in files):
    raise ValueError("SOG contains unexpected asset references")
with ThreadPoolExecutor(max_workers=3) as pool:
    rows = list(pool.map(lambda name: fetch(name, 100 * 1024**2), files))
report = {
    "schema_version": "missioncore.ai-polygon-source-archive/v1",
    "source_url": args.source_url,
    "metadata_url": args.metadata_url,
    "author": args.author,
    "license": args.license,
    "captured_at": datetime.now(UTC).isoformat(),
    "elapsed_ns": time.monotonic_ns() - started,
    "splat_count": metadata["count"],
    "files": [meta, *rows],
}
(args.output / "source-manifest.json").write_text(json.dumps(report, indent=2))
print(json.dumps({"count": report["splat_count"], "bytes": sum(r["byte_length"] for r in rows)}))