#!/usr/bin/env python3
"""Open viser viewer with all PLYs from one AUV.

Usage:
    viser_auv.py --ply-dir /path/to/auv/ply --port 9210
"""
from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import numpy as np


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ply-dir", required=True)
    ap.add_argument("--port", type=int, default=9210)
    ap.add_argument("--point-size", type=float, default=0.01)
    ap.add_argument("--max-points-per-ply", type=int, default=1_500_000)
    args = ap.parse_args()

    try:
        import open3d as o3d
        import viser
    except ImportError as e:
        sys.exit(f"missing dep: {e}")

    ply_dir = Path(args.ply_dir)
    plys = sorted(ply_dir.glob("**/*.ply"))
    print(f"Found {len(plys)} PLY files in {ply_dir}", flush=True)
    if not plys:
        sys.exit("no PLY found")

    server = viser.ViserServer(host="0.0.0.0", port=args.port)
    palette = [
        (1.0, 0.30, 0.30), (0.30, 1.0, 0.30), (0.30, 0.55, 1.0),
        (1.0, 0.85, 0.20), (1.0, 0.30, 1.0), (0.30, 1.0, 1.0),
        (1.0, 0.55, 0.20), (0.55, 0.30, 1.0),
    ]

    for i, p in enumerate(plys):
        pcd = o3d.io.read_point_cloud(str(p))
        pts = np.asarray(pcd.points, dtype=np.float32)
        if len(pts) == 0:
            print(f"  ! {p.name}: empty", flush=True)
            continue
        if pcd.has_colors():
            cols = np.asarray(pcd.colors, dtype=np.float32)
        else:
            cols = np.tile(palette[i % len(palette)], (len(pts), 1)).astype(np.float32)
        if len(pts) > args.max_points_per_ply:
            idx = np.random.choice(len(pts), args.max_points_per_ply, replace=False)
            pts = pts[idx]
            cols = cols[idx]
        # viser wants uint8 colors
        cols_u8 = (cols * 255).clip(0, 255).astype(np.uint8)
        name = f"/{p.parent.name}_{p.stem}"
        server.scene.add_point_cloud(
            name=name, points=pts, colors=cols_u8, point_size=args.point_size
        )
        print(f"  + {p.name}: {len(pts):,} pts", flush=True)

    print(f"Viser ready on port {args.port}", flush=True)
    while True:
        time.sleep(60)


if __name__ == "__main__":
    main()