#!/usr/bin/env python3
"""Refresh metadata.csv status columns for Edit3D-Bench TRELLIS preprocessing."""

import argparse
from pathlib import Path

import pandas as pd


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output_dir", required=True)
    parser.add_argument("--feat_model", default="dinov2_vitl14_reg")
    parser.add_argument("--ss_latent_name", default="ss_enc_conv3d_16l8_fp16")
    parser.add_argument("--latent_name", default="dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    metadata_path = output_dir / "metadata.csv"
    df = pd.read_csv(metadata_path)

    rendered = []
    voxelized = []
    num_voxels = []
    feature_col = f"feature_{args.feat_model}"
    ss_col = f"ss_latent_{args.ss_latent_name}"
    latent_col = f"latent_{args.latent_name}"
    feature_ok = []
    ss_ok = []
    latent_ok = []

    for row in df.to_dict("records"):
        sha = row["sha256"]
        render_dir = output_dir / "renders" / sha
        voxel_path = output_dir / "voxels" / f"{sha}.ply"
        feature_path = output_dir / "features" / args.feat_model / f"{sha}.npz"
        ss_path = output_dir / "ss_latents" / args.ss_latent_name / f"{sha}.npz"
        latent_path = output_dir / "latents" / args.latent_name / f"{sha}.npz"
        rendered.append((render_dir / "transforms.json").exists() and (render_dir / "mesh.ply").exists())
        voxelized.append(voxel_path.exists())
        feature_ok.append(feature_path.exists())
        ss_ok.append(ss_path.exists())
        latent_ok.append(latent_path.exists())
        if voxel_path.exists():
            try:
                import utils3d

                num_voxels.append(int(len(utils3d.io.read_ply(str(voxel_path))[0])))
            except Exception:
                num_voxels.append(0)
        else:
            num_voxels.append(0)

    df["rendered"] = rendered
    df["voxelized"] = voxelized
    df["num_voxels"] = num_voxels
    df[feature_col] = feature_ok
    df[ss_col] = ss_ok
    df[latent_col] = latent_ok
    df.to_csv(metadata_path, index=False)

    print(f"metadata: {metadata_path}")
    print(f"rendered {sum(rendered)}/{len(df)}")
    print(f"voxelized {sum(voxelized)}/{len(df)}")
    print(f"{feature_col} {sum(feature_ok)}/{len(df)}")
    print(f"{ss_col} {sum(ss_ok)}/{len(df)}")
    print(f"{latent_col} {sum(latent_ok)}/{len(df)}")


if __name__ == "__main__":
    main()