#!/usr/bin/env python3
"""Check integrity of processed jobs (PLY + poses present)."""
import argparse
import glob
import json
import os
from pathlib import Path
from typing import Any

REQUIRED_FILES = ["reconstruction.ply", "lingbot_poses.npz"]
OPTIONAL_FILES = ["model_decimated.ply"]


def check_job(job_id: int, frames_base: str = "/root/cosma-qc-frames") -> dict[str, Any]:
    job_dir = Path(frames_base) / f"job_{job_id}"
    if not job_dir.exists():
        return {"job_id": job_id, "status": "missing", "missing": [], "details": {}}

    missing = [f for f in REQUIRED_FILES if not (job_dir / f).exists()]
    details: dict[str, Any] = {}

    ply = job_dir / "reconstruction.ply"
    if ply.exists():
        details["ply_size_gb"] = round(ply.stat().st_size / 1e9, 2)

    poses = job_dir / "lingbot_poses.npz"
    if poses.exists():
        try:
            import numpy as np
            d = np.load(str(poses))
            n = d["poses"].shape[0] if "poses" in d else 0
            details["n_poses"] = n
        except Exception as e:
            details["poses_error"] = str(e)

    decimated = job_dir / "model_decimated.ply"
    details["decimated"] = decimated.exists()

    return {
        "job_id": job_id,
        "status": "ok" if not missing else "incomplete",
        "missing": missing,
        "details": details,
    }


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("job_ids", nargs="*", type=int)
    p.add_argument("--frames-base", default="/root/cosma-qc-frames")
    p.add_argument("--all", action="store_true", help="Check all job dirs")
    args = p.parse_args()

    base = Path(args.frames_base)
    if args.all:
        ids = sorted(
            int(d.name.replace("job_", ""))
            for d in base.iterdir()
            if d.is_dir() and d.name.startswith("job_") and d.name[4:].isdigit()
        )
    else:
        ids = args.job_ids

    results = [check_job(jid, args.frames_base) for jid in ids]
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()