"""
Batch fix USD metadata for IS 5.0.0 compatibility.

IS 4.5.0 auto-added 'Rotate:unitsResolve' and 'Scale:unitsResolve' xformOps
to compensate for metersPerUnit/upAxis differences. IS 5.0.0 no longer does this,
causing RigidObject physics to break (objects fly away).

Fix: Change metersPerUnit to 1.0 and upAxis to Z in USD metadata.
Vertex data stays unchanged (already in correct scale for the scene).

Usage:
    # Scan only (dry run):
    python migerate/fix_usd_metadata.py --root workflows/simbox/example_assets --dry-run

    # Fix all:
    python migerate/fix_usd_metadata.py --root workflows/simbox/example_assets

    # Fix specific pattern:
    python migerate/fix_usd_metadata.py --root workflows/simbox/example_assets --pattern "Aligned_obj.usd"

    # Skip specific directories (e.g. robot models, curobo assets):
    python migerate/fix_usd_metadata.py --root . --skip curobo,robot

    # Restore from backups:
    python migerate/fix_usd_metadata.py --root workflows/simbox/example_assets --restore
"""
import argparse
import glob
import os
import shutil
import sys

from pxr import Usd, UsdGeom


def fix_usd_files(root: str, pattern: str = "*.usd", target_mpu: float = 1.0,
                   target_up: str = "Z", dry_run: bool = False, skip: list = None):
    usd_files = glob.glob(os.path.join(root, "**", pattern), recursive=True)
    usd_files = [f for f in usd_files if not f.endswith(".bak")]
    usd_files.sort()

    skip = skip or []

    print(f"Scanning {len(usd_files)} USD files under: {root}")
    if dry_run:
        print("[DRY RUN] No files will be modified.\n")
    print()

    fixed = 0
    skipped = 0
    skipped_dir = 0
    errors = 0

    for fpath in usd_files:
        rel = os.path.relpath(fpath, root)

        # Check skip patterns
        if any(s in rel for s in skip):
            print(f"  SKIP (dir filter): {rel}")
            skipped_dir += 1
            continue

        try:
            stage = Usd.Stage.Open(fpath)
            if stage is None:
                print(f"  ERROR: Cannot open {rel}")
                errors += 1
                continue

            mpu = UsdGeom.GetStageMetersPerUnit(stage)
            up = UsdGeom.GetStageUpAxis(stage)

            needs_fix = (mpu != target_mpu) or (up != target_up)
            if not needs_fix:
                print(f"  OK: {rel} (mpu={mpu}, up={up})")
                skipped += 1
                continue

            if dry_run:
                print(f"  WOULD FIX: {rel} (mpu={mpu}, up={up} -> mpu={target_mpu}, up={target_up})")
                fixed += 1
                continue

            # Backup
            bak = fpath + ".bak"
            if not os.path.exists(bak):
                shutil.copy2(fpath, bak)

            # Fix
            UsdGeom.SetStageMetersPerUnit(stage, target_mpu)
            UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z if target_up == "Z" else UsdGeom.Tokens.y)
            stage.GetRootLayer().Save()

            print(f"  FIXED: {rel} (mpu={mpu}/{up} -> {target_mpu}/{target_up})")
            fixed += 1

        except Exception as e:
            print(f"  ERROR: {rel} -> {e}")
            errors += 1

    print(f"\nDone: {fixed} {'would fix' if dry_run else 'fixed'}, {skipped} already OK, "
          f"{skipped_dir} skipped (dir filter), {errors} errors.")

    if not dry_run and fixed > 0:
        print("Backups saved as *.bak.")
        print("To restore: python migerate/fix_usd_metadata.py --root <DIR> --restore")

    return fixed


def restore_backups(root: str, pattern: str = "*.usd"):
    bak_files = glob.glob(os.path.join(root, "**", pattern + ".bak"), recursive=True)
    bak_files.sort()

    if not bak_files:
        print(f"No .bak files found under: {root}")
        return

    print(f"Found {len(bak_files)} backup files. Restoring...")
    for bak in bak_files:
        original = bak[:-4]  # remove .bak
        shutil.copy2(bak, original)
        print(f"  Restored: {os.path.relpath(original, root)}")
    print(f"\nDone: {len(bak_files)} files restored.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Batch fix USD metersPerUnit/upAxis metadata for IS 5.0.0 compatibility",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Dry run on example assets:
  python migerate/fix_usd_metadata.py --root workflows/simbox/example_assets --dry-run

  # Fix all Aligned_obj.usd in example assets:
  python migerate/fix_usd_metadata.py --root workflows/simbox/example_assets --pattern "Aligned_obj.usd"

  # Fix everything except robot/curobo dirs:
  python migerate/fix_usd_metadata.py --root . --skip curobo,robot

  # Restore all backups:
  python migerate/fix_usd_metadata.py --root workflows/simbox/example_assets --restore
""")
    parser.add_argument("--root", default="workflows/simbox/example_assets",
                        help="Root directory to scan (default: workflows/simbox/example_assets)")
    parser.add_argument("--pattern", default="*.usd", help="Filename glob pattern (default: *.usd)")
    parser.add_argument("--target-mpu", type=float, default=1.0, help="Target metersPerUnit (default: 1.0)")
    parser.add_argument("--target-up", choices=["Y", "Z"], default="Z", help="Target upAxis (default: Z)")
    parser.add_argument("--dry-run", action="store_true", help="Scan and report only, do not modify files")
    parser.add_argument("--skip", default="", help="Comma-separated dir substrings to skip (e.g. 'curobo,robot')")
    parser.add_argument("--restore", action="store_true", help="Restore all .bak files to originals")
    args = parser.parse_args()

    if args.restore:
        restore_backups(args.root, args.pattern)
    else:
        skip_list = [s.strip() for s in args.skip.split(",") if s.strip()]
        fix_usd_files(args.root, args.pattern, args.target_mpu, args.target_up, args.dry_run, skip_list)