#!/usr/bin/env python3
"""Decimate PLY and SCP to cosma-vm after a job completes."""
import argparse
import subprocess
import sys
from pathlib import Path

COSMA_VM = "cosma@192.168.0.83"
COSMA_DATA = "/data/cosma"
MAX_PTS = 300_000


def find_ply(frames_dir: Path) -> Path | None:
    for candidate in ["model.ply", "output.ply", "reconstruction.ply"]:
        p = frames_dir / candidate
        if p.exists():
            return p
    plys = list(frames_dir.glob("*.ply"))
    return plys[0] if plys else None


def decimate_ply(src: str, dst: str, max_pts: int = MAX_PTS) -> None:
    import open3d as o3d
    import numpy as np
    src_path = Path(src)
    if not src_path.exists():
        raise FileNotFoundError(src)
    pcd = o3d.io.read_point_cloud(str(src_path))
    n = len(pcd.points)
    if n > max_pts:
        vol = float(np.prod(pcd.get_max_bound() - pcd.get_min_bound()))
        vox = max((vol / max_pts) ** (1 / 3), 0.02)
        pcd = pcd.voxel_down_sample(vox)
    o3d.io.write_point_cloud(dst, pcd)
    print(f"Decimated {n} -> {len(pcd.points)} pts -> {dst}", flush=True)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("job_id", type=int)
    p.add_argument("--frames-dir", required=True)
    p.add_argument("--cosma-vm", default=COSMA_VM)
    p.add_argument("--cosma-data", default=COSMA_DATA)
    args = p.parse_args()

    frames_dir = Path(args.frames_dir)
    ply_src = find_ply(frames_dir)
    if ply_src is None:
        print(f"No PLY found in {frames_dir}", flush=True)
        sys.exit(0)

    ply_dec = frames_dir / "model_decimated.ply"
    decimate_ply(str(ply_src), str(ply_dec))

    remote_dir = args.cosma_data
    subprocess.run(["ssh", args.cosma_vm, f"mkdir -p {remote_dir}"], check=True)
    dst = f"job_{args.job_id}_decimated.ply"
    subprocess.run(["scp", str(ply_dec), f"{args.cosma_vm}:{remote_dir}/{dst}"], check=True)
    print(f"SCP done -> {args.cosma_vm}:{remote_dir}/{dst}", flush=True)

    poses_src = frames_dir / "lingbot_poses.npz"
    if poses_src.exists():
        poses_dst = f"job_{args.job_id}_poses.npz"
        subprocess.run(["scp", str(poses_src), f"{args.cosma_vm}:{remote_dir}/{poses_dst}"], check=True)
        print(f"SCP poses -> {args.cosma_vm}:{remote_dir}/{poses_dst}", flush=True)


if __name__ == "__main__":
    main()