"""Prepare paired, canonical rendered images and coarse Gaussian targets.

These are mesh-derived proxy targets, NOT pretrained-teacher outputs or photos.
Only public, explicitly allowed assets are downloaded. Source meshes stay local.
"""
import argparse
import gzip
import hashlib
import json
import random
import shutil
import time
import urllib.error
import urllib.request
from pathlib import Path, PurePosixPath

import numpy as np
import pyarrow.parquet as pq
import trimesh
from huggingface_hub import hf_hub_download

from image3d_student.data import cache_sample, splat_targets
from image3d_student.splats import Splats, preview
from image3d_student.engine import split_rows

SCENARIO = ("scenario-labs/3d-collection", "55b4c002b4c1529fc48da3dff98516c3069bfa3f")
OBJAVERSE = ("allenai/objaverse", "556637099bf4fa79ea7b239d0d9c328b8a2e9ac8")


def digest(path):
    h = hashlib.sha256()
    with Path(path).open("rb") as f:
        for block in iter(lambda: f.read(1024 * 1024), b""):
            h.update(block)
    return h.hexdigest()


class Downloader:
    def __init__(self, root, max_file_mb, max_total_mb, reserve_gb):
        self.root = Path(root)
        self.root.mkdir(parents=True, exist_ok=True)
        self.max_file = int(max_file_mb * 1024**2)
        self.max_total = int(max_total_mb * 1024**2)
        self.reserve = int(reserve_gb * 1024**3)
        self.total = 0

    def get(self, repo, revision, filename, expected=None):
        parts = PurePosixPath(filename)
        if parts.is_absolute() or ".." in parts.parts or not filename.endswith(".glb"):
            raise ValueError("Only relative GLB paths are allowed")
        out = self.root / repo.replace("/", "--") / revision / filename
        if out.exists():
            if expected and digest(out) != expected:
                raise ValueError("Cached source checksum mismatch")
            if out.stat().st_size > self.max_file:
                raise ValueError("Cached asset exceeds per-file limit")
            return out
        if self.total >= self.max_total:
            raise RuntimeError("Download budget exhausted")
        out.parent.mkdir(parents=True, exist_ok=True)
        url = f"https://huggingface.co/datasets/{repo}/resolve/{revision}/{filename}"
        for attempt in range(4):
            tmp = out.with_suffix(".part")
            try:
                # Public URL only: credentials are never forwarded to asset hosts.
                with urllib.request.urlopen(url, timeout=60) as response, tmp.open("wb") as f:
                    length = int(response.headers.get("Content-Length", 0))
                    if length > self.max_file:
                        raise ValueError("Asset exceeds per-file limit")
                    count = 0
                    while block := response.read(1024 * 1024):
                        count += len(block)
                        self.total += len(block)
                        if count > self.max_file:
                            raise ValueError("Asset exceeds per-file limit")
                        if self.total > self.max_total:
                            raise RuntimeError("Download budget exhausted")
                        if shutil.disk_usage(self.root).free < self.reserve + len(block):
                            raise RuntimeError("Free disk space below reserve")
                        f.write(block)
                if expected and digest(tmp) != expected:
                    raise ValueError("Source checksum mismatch")
                tmp.replace(out)
                return out
            except urllib.error.HTTPError as exc:
                if exc.code not in (429, 500, 502, 503, 504) or attempt == 3:
                    raise
                time.sleep(min(2**attempt, 8))
            finally:
                tmp.unlink(missing_ok=True)
        raise RuntimeError("Download retries exhausted")


def scenario_candidates(cache, seed):
    repo, revision = SCENARIO
    path = hf_hub_download(repo, "data/train-00000-of-00001.parquet", repo_type="dataset",
                           revision=revision, local_dir=Path(cache) / "metadata/scenario")
    rows = pq.read_table(path, columns=["id", "license", "model", "inputs", "outputs"]).to_pylist()
    # Connected components keep alternate meshes generated from shared inputs together.
    parent = {}

    def find(x):
        parent.setdefault(x, x)
        if parent[x] != x:
            parent[x] = find(parent[x])
        return parent[x]

    candidates = []
    for row in rows:
        if row["license"] != "cc-by-4.0" or "hunyuan" in (row["model"] or "").lower():
            continue
        assets = (row.get("inputs") or {}).get("assets") or []
        keys = ["input:" + (a.get("sha256") or a["path"]) for a in assets if a.get("sha256") or a.get("path")]
        keys += ["output:" + o["sha256"] for o in row["outputs"] if o.get("sha256")]
        keys += [row["id"]]
        for key in keys[1:]:
            parent[find(key)] = find(keys[0])
        for output in row["outputs"]:
            if output.get("format") == "glb":
                candidates.append({"repo": repo, "revision": revision, "path": output["path"],
                    "sha256": output.get("sha256"), "asset_id": row["id"], "license": "cc-by-4.0",
                    "author": "Scenario", "title": row["id"], "source_model": row["model"],
                    "source_url": f"https://huggingface.co/datasets/{repo}", "group_key": keys[0]})
                break
    for row in candidates:
        row["object_id"] = "scenario-" + hashlib.sha256(find(row.pop("group_key")).encode()).hexdigest()[:20]
    candidates.sort(key=lambda r: r["asset_id"])
    random.Random(seed).shuffle(candidates)
    return candidates


def objaverse_candidates(cache, seed, shards):
    repo, revision = OBJAVERSE
    candidates = []
    for shard in range(shards):
        part = f"000-{shard:03d}"
        path = hf_hub_download(repo, f"metadata/{part}.json.gz", repo_type="dataset", revision=revision,
                               local_dir=Path(cache) / "metadata/objaverse")
        with gzip.open(path) as f:
            records = json.load(f)
        for uid, row in records.items():
            # Start with the more conservative CC0 subset; no NC/SA/unknown assets.
            if row.get("license") != "cc0" or row.get("isAgeRestricted") or row.get("animationCount", 0):
                continue
            candidates.append({"repo": repo, "revision": revision, "path": f"glbs/{part}/{uid}.glb",
                "asset_id": uid, "object_id": "objaverse-" + uid, "license": "cc0-1.0",
                "author": row.get("user", {}).get("displayName", "Unknown"), "title": row.get("name", uid),
                "source_url": row.get("viewerUrl", f"https://sketchfab.com/3d-models/{uid}")})
    candidates.sort(key=lambda r: r["asset_id"])
    random.Random(seed).shuffle(candidates)
    return candidates


def mesh_splats(path, count, grid, seed):
    scene = trimesh.load_scene(path, process=False)
    meshes = []
    for node in scene.graph.nodes_geometry:
        matrix, name = scene.graph[node]
        mesh = scene.geometry[name].copy()
        if not isinstance(mesh, trimesh.Trimesh) or not len(mesh.faces):
            continue
        mesh.apply_transform(matrix)
        meshes.append(mesh)
    if not meshes or sum(len(m.faces) for m in meshes) > 500_000:
        raise ValueError("Missing mesh or more than 500,000 triangles")
    bounds = np.array([m.bounds for m in meshes])
    low, high = bounds[:, 0].min(axis=0), bounds[:, 1].max(axis=0)
    extent = float((high - low).max())
    if not np.isfinite(bounds).all() or extent < 1e-8:
        raise ValueError("Degenerate mesh")
    center = (low + high) / 2
    factor = .9 / extent
    areas = np.array([m.area for m in meshes])
    if not np.isfinite(areas).all() or areas.sum() <= 0:
        raise ValueError("Invalid surface area")
    counts = np.random.default_rng(seed).multinomial(count, areas / areas.sum())
    points, colours = [], []
    for mesh, n in zip(meshes, counts):
        if not n:
            continue
        xyz, faces = trimesh.sample.sample_surface(mesh, int(n), seed=seed)
        bary = trimesh.triangles.points_to_barycentric(mesh.triangles[faces], xyz)
        vis = mesh.visual
        if vis.kind == "texture" and vis.uv is not None:
            uv = (vis.uv[mesh.faces[faces]] * bary[:, :, None]).sum(axis=1)
            sampled = np.asarray(vis.material.to_color(uv))
            if sampled.ndim == 0:
                raise ValueError("Material has no usable texture or base colour")
            # Untextured PBR materials can return one constant RGBA value.
            if sampled.ndim == 1:
                sampled = np.broadcast_to(sampled, (n, len(sampled)))
            colour = sampled[:, :3] / 255.
        elif vis.kind == "vertex":
            colour = (vis.vertex_colors[mesh.faces[faces], :3] * bary[:, :, None]).sum(axis=1) / 255.
        elif vis.kind == "face":
            colour = vis.face_colors[faces, :3] / 255.
        else:
            colour = np.full((n, 3), .65)
        points.append((xyz - center) * factor)
        colours.append(colour)
    xyz = np.concatenate(points).astype(np.float32)
    rgb = np.concatenate(colours).astype(np.float32)
    n = len(xyz)
    splats = Splats(xyz, rgb, np.full((n, 3), .6 / grid),
                    np.tile([1, 0, 0, 0], (n, 1)), np.full((n, 1), .9))
    transform = np.eye(4)
    transform[:3, :3] *= factor
    transform[:3, 3] = -center * factor
    return splats, transform.tolist()


def save_progress(out, args, rows, provenance, failures):
    state = {"version": 1, "settings": vars(args), "rows": rows,
             "sources": provenance, "skipped": failures}
    temporary = out / "progress.json.tmp"
    temporary.write_text(json.dumps(state))
    temporary.replace(out / "progress.json")


def check_resume_settings(previous, args):
    for key in ("grid", "views", "surface_points", "seed", "scenario", "objaverse"):
        if previous[key] != getattr(args, key):
            raise ValueError(f"Cannot change {key} while resuming; use a new dataset directory")


def recover_legacy(out, args):
    """Recover the old writer's completed files without rendering or downloading.

    Old runs saved sources and arrays, but lost manifest rows on insufficient
    candidates. Recreate those rows and verify targets against cached GLBs.
    """
    previous = json.loads((out / "preparation.json").read_text())["settings"]
    check_resume_settings(previous, args)
    provenance = json.loads((out / "sources.json").read_text())
    failures = json.loads((out / "skipped.json").read_text())
    rows = []
    for index, item in enumerate(provenance):
        path = Path(args.cache) / item["repo"].replace("/", "--") / item["revision"] / item["path"]
        if not path.is_file() or digest(path) != item["sha256"]:
            raise ValueError(f"Missing or corrupt cached source for {item['asset_id']}")
        splats, transform = mesh_splats(path, args.surface_points, args.grid,
                                        int(item["sha256"][:8], 16) ^ args.seed)
        stem = hashlib.sha256((item["repo"] + item["asset_id"]).encode()).hexdigest()[:20]
        for view in range(args.views):
            sample_id = f"{stem}-v{view}"
            image, target = f"images/{sample_id}.png", f"targets/{sample_id}.npz"
            yaw = 2 * np.pi * view / args.views
            c, s = np.cos(yaw), np.sin(yaw)
            rotation = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], np.float32)
            rotated = Splats(splats.xyz @ rotation.T, splats.rgb, splats.scale, splats.quat, splats.opacity)
            expected = splat_targets(rotated, args.grid)
            with np.load(out / target, allow_pickle=False) as cached:
                if set(cached.files) != set(expected):
                    raise ValueError(f"Target fields differ for {sample_id}")
                for key, value in expected.items():
                    if not np.array_equal(cached[key], value):
                        raise ValueError(f"Recovered target differs for {sample_id}: {key}")
            from PIL import Image
            with Image.open(out / image) as check:
                check.verify()
            rows.append({"id": sample_id, "image": image, "target": target,
                "source": "mesh-proxy:" + item["repo"], "grid": args.grid,
                "teachers": {"appearance": "mesh-proxy:" + item["repo"], "coarse": None, "fine": None},
                "teacher_gaussians": len(splats.xyz), "retained_gaussians": int(expected["occ"].sum()),
                "geometry_transforms": {"coarse": None, "fine": None},
                "image_sha256": digest(out / image), "target_sha256": digest(out / target),
                "metadata": {**item, "normalization_matrix": transform, "yaw_radians": float(yaw),
                    "render": "orthographic Gaussian proxy, camera +z",
                    "target": "isotropic surface-sampled mesh proxy; not a pretrained teacher"}})
        if (index + 1) % 10 == 0 or index + 1 == len(provenance):
            print(json.dumps({"recovered_objects": index + 1, "existing_objects": len(provenance)}), flush=True)
    save_progress(out, args, rows, provenance, failures)
    return rows, provenance, failures


def load_progress(out, args):
    if not (out / "progress.json").exists():
        return recover_legacy(out, args)
    state = json.loads((out / "progress.json").read_text())
    if state.get("version") != 1:
        raise ValueError("Unsupported preparation progress version")
    check_resume_settings(state["settings"], args)
    for row in state["rows"]:
        for key in ("image", "target"):
            path = (out / row[key]).resolve()
            if not path.is_relative_to(out.resolve()) or digest(path) != row[key + "_sha256"]:
                raise ValueError(f"Prepared file checksum mismatch: {row['id']} {key}")
    return state["rows"], state["sources"], state["skipped"]


def prune_invisible(out, rows, provenance, failures):
    from PIL import Image
    bad = set()
    for row in rows:
        with Image.open(out / row["image"]) as image:
            if image.convert("RGB").getbbox() is None:
                bad.add((row["metadata"]["repo"], row["metadata"]["asset_id"]))
    for repo, asset in sorted(bad):
        failures.append({"repo": repo, "asset": asset,
                         "error": "Rendered view has no visible pixels", "terminal": True})
    if bad:
        print(json.dumps({"excluded_invisible_objects": len(bad), "existing_files_preserved": True}), flush=True)
    return ([r for r in rows if (r["metadata"]["repo"], r["metadata"]["asset_id"]) not in bad],
            [p for p in provenance if (p["repo"], p["asset_id"]) not in bad])


def group_duplicates(rows, provenance):
    """Link entire objects when any views or target files are byte-identical."""
    parent = {}
    def find(x):
        parent.setdefault(x, x)
        if parent[x] != x:
            parent[x] = find(parent[x])
        return parent[x]
    hashes = {}
    for row in rows:
        meta = row["metadata"]
        original = meta.setdefault("source_object_id", meta["object_id"])
        find(original)
        for field in ("image_sha256", "target_sha256"):
            key = (field, row[field])
            if key in hashes:
                a, b = find(original), find(hashes[key])
                parent[max(a, b)] = min(a, b)
            else:
                hashes[key] = original
    for row in rows:
        row["metadata"]["object_id"] = find(row["metadata"]["source_object_id"])
    for item in provenance:
        original = item.setdefault("source_object_id", item["object_id"])
        item["object_id"] = find(original)


def prepare(args):
    out = Path(args.out)
    occupied = out.exists() and any(out.iterdir())
    if occupied and not getattr(args, "resume", False):
        raise ValueError("Output is not empty. Use --resume to retain completed objects.")
    if args.scenario + args.objaverse < 4:
        raise ValueError("Request at least four objects")
    out.mkdir(parents=True, exist_ok=True)
    if occupied:
        all_rows, provenance, failures = load_progress(out, args)
    else:
        all_rows, provenance, failures = [], [], []
        save_progress(out, args, all_rows, provenance, failures)
    all_rows, provenance = prune_invisible(out, all_rows, provenance, failures)
    save_progress(out, args, all_rows, provenance, failures)
    seen = {item["sha256"] for item in provenance}
    completed = {(item["repo"], item["asset_id"]) for item in provenance}
    rejected = {(f["repo"], f["asset"]) for f in failures if f.get("terminal")}
    counts = {}
    fetch = Downloader(args.cache, args.max_file_mb, args.max_download_mb, args.reserve_gb)
    sources = []
    if args.scenario:
        ready = sum(item["repo"] == SCENARIO[0] for item in provenance)
        candidates = scenario_candidates(args.cache, args.seed) if ready < args.scenario else []
        sources.append((SCENARIO[0], candidates, args.scenario, ready))
    if args.objaverse:
        ready = sum(item["repo"] == OBJAVERSE[0] for item in provenance)
        candidates = objaverse_candidates(args.cache, args.seed, args.objaverse_shards) if ready < args.objaverse else []
        sources.append((OBJAVERSE[0], candidates, args.objaverse, ready))
    for repo, candidates, wanted, accepted in sources:
        for item in candidates:
            if accepted >= wanted:
                break
            if (item["repo"], item["asset_id"]) in completed | rejected:
                continue
            staging = out / "_staging"
            try:
                path = fetch.get(item["repo"], item["revision"], item["path"], item.get("sha256"))
                checksum = digest(path)
                if checksum in seen:
                    continue
                item["sha256"] = checksum
                item_seed = int(checksum[:8], 16) ^ args.seed
                splats, transform = mesh_splats(path, args.surface_points, args.grid, item_seed)
                stem = hashlib.sha256((item["repo"] + item["asset_id"]).encode()).hexdigest()[:20]
                staging.mkdir(parents=True, exist_ok=True)
                rows = []
                for view in range(args.views):
                    yaw = 2 * np.pi * view / args.views
                    c, s = np.cos(yaw), np.sin(yaw)
                    rotation = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], np.float32)
                    rotated = Splats(splats.xyz @ rotation.T, splats.rgb, splats.scale, splats.quat, splats.opacity)
                    image = staging / "input.png"
                    rendered = preview(rotated, size=128)
                    if rendered.getbbox() is None:
                        raise ValueError("Rendered view has no visible pixels")
                    rendered.save(image)
                    metadata = {**item, "object_id": item["object_id"], "normalization_matrix": transform,
                                "yaw_radians": float(yaw), "render": "orthographic Gaussian proxy, camera +z",
                                "target": "isotropic surface-sampled mesh proxy; not a pretrained teacher"}
                    rows.append(cache_sample(staging, f"{stem}-v{view}", image, rotated, args.grid,
                                             source="mesh-proxy:" + item["repo"], extra=metadata))
                for folder in ("images", "targets"):
                    (out / folder).mkdir(exist_ok=True)
                    for file in (staging / folder).iterdir():
                        file.replace(out / folder / file.name)
                all_rows.extend(rows)
                seen.add(checksum)
                provenance.append(item)
                completed.add((item["repo"], item["asset_id"]))
                accepted += 1
                save_progress(out, args, all_rows, provenance, failures)
                print(json.dumps({"repo": item["repo"], "accepted": accepted, "requested": wanted,
                                  "asset": item["asset_id"]}), flush=True)
            except RuntimeError:
                raise  # resource limits are a hard stop
            except Exception as exc:
                failures.append({"repo": item["repo"], "asset": item["asset_id"],
                                 "error": type(exc).__name__ + ": " + str(exc)[:200],
                                 "terminal": "no visible pixels" in str(exc)})
                print(json.dumps({"skipped": item["asset_id"], "reason": failures[-1]["error"]}), flush=True)
            finally:
                shutil.rmtree(staging, ignore_errors=True)
        counts[repo] = {"requested": wanted, "prepared": accepted}
    group_duplicates(all_rows, provenance)
    save_progress(out, args, all_rows, provenance, failures)
    (out / "sources.json").write_text(json.dumps(provenance, indent=2))
    (out / "skipped.json").write_text(json.dumps(failures, indent=2))
    report = {"counts": counts, "samples": len(all_rows), "objects": len(provenance),
              "distinct_groups": len({r["metadata"]["object_id"] for r in all_rows}),
              "settings": vars(args), "rendered_proxy_data": True,
              "excluded": ["scenario-labs/3d-collection-hunyuan", "3dlg-hcvc/MultiScan",
                           "huanngzh/3D-Front", "facebook/show3d-dataset", "Voxel51/gaussian_splatting"]}
    (out / "preparation.json").write_text(json.dumps(report, indent=2))
    if any(v["prepared"] < v["requested"] for v in counts.values()):
        raise ValueError("Not enough eligible assets. Completed objects are saved. Rerun with --resume and a larger --objaverse-shards value.")
    if report["distinct_groups"] < 4:
        raise ValueError("Not enough independent object groups")
    # The test partition never enters training or checkpoint selection.
    development, test = split_rows(all_rows, .15, args.seed + 1)
    for name, rows in [("manifest.jsonl", development), ("test-manifest.jsonl", test)]:
        (out / name).write_text("".join(json.dumps(r) + "\n" for r in rows))
    print(json.dumps(report, indent=2))


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--out", required=True)
    p.add_argument("--cache", default="data/downloads")
    p.add_argument("--scenario", type=int, default=64)
    p.add_argument("--objaverse", type=int, default=32)
    p.add_argument("--objaverse-shards", type=int, default=8)
    p.add_argument("--resume", action="store_true", help="Reuse and verify completed objects, including older failed runs")
    p.add_argument("--grid", type=int, choices=[8, 16, 32], default=16)
    p.add_argument("--views", type=int, default=4)
    p.add_argument("--surface-points", type=int, default=4096)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--max-file-mb", type=int, default=16)
    p.add_argument("--max-download-mb", type=int, default=2048)
    p.add_argument("--reserve-gb", type=float, default=2)
    args = p.parse_args()
    if min(args.scenario, args.objaverse) < 0 or not 1 <= args.objaverse_shards <= 160:
        p.error("Invalid object counts or shard count")
    if min(args.views, args.surface_points, args.max_file_mb, args.max_download_mb, args.reserve_gb) <= 0:
        p.error("Counts and limits must be positive")
    prepare(args)


if __name__ == "__main__":
    main()