#!/usr/bin/env python3
"""Migrate World-UAV-aug from flat prefix layout to directory-based layout.

Old layout (prefix-based):
    World-UAV-aug/Rot/scene/DB/img/crop_12_4_depth.png
    World-UAV-aug/Rot/scene/DB/img/crop_12_4_edge.png
    World-UAV-aug/Rot/scene/DB/img/crop_12_4_segm.png
    World-UAV-aug/Rot/scene/DB/img/crop_12_4_chm.png
    World-UAV-aug/Rot/scene/DB/img/crop_12_4_depth.npy
    ...

New layout (directory-based):
    World-UAV-aug/depth/Rot/scene/DB/img/crop_12_4.png
    World-UAV-aug/edge/Rot/scene/DB/img/crop_12_4.png
    World-UAV-aug/segm/Rot/scene/DB/img/crop_12_4.png
    World-UAV-aug/chm/Rot/scene/DB/img/crop_12_4.png
    World-UAV-aug/npy/depth/Rot/scene/DB/img/crop_12_4.npy
    ...

Usage:
    python scripts/migrate_layout.py /mnt/data1tb/cvgl_datasets/World-UAV-aug
    python scripts/migrate_layout.py /mnt/data1tb/cvgl_datasets/World-UAV-aug --dry-run
"""

from __future__ import annotations

import argparse
import os
import shutil
from pathlib import Path

# Suffix → (modality, extension)
_SUFFIX_MAP = {
    "_depth.png": ("depth", ".png"),
    "_depth.npy": ("depth", ".npy"),
    "_edge.png":  ("edge",  ".png"),
    "_edge.npy":  ("edge",  ".npy"),
    "_segm.png":  ("segm",  ".png"),
    "_segm.npy":  ("segm",  ".npy"),
    "_chm.png":   ("chm",   ".png"),
    "_chm.npy":   ("chm",   ".npy"),
}

# Top-level dirs to skip (they belong to the new layout).
_SKIP_DIRS = {"depth", "edge", "segm", "chm", "npy", "safetensors"}


def find_old_files(root: Path) -> list[tuple[Path, str, str, str]]:
    """Find all files with old prefix-based naming.

    Returns list of (old_path, rel_parent, stem, suffix_key).
    """
    results = []
    for p in sorted(root.rglob("*")):
        if not p.is_file():
            continue
        # Skip files already in new-layout dirs.
        try:
            rel = p.relative_to(root)
        except ValueError:
            continue
        if rel.parts and rel.parts[0] in _SKIP_DIRS:
            continue

        name = p.name
        for suffix_key, (modality, ext) in _SUFFIX_MAP.items():
            if name.endswith(suffix_key):
                stem = name[: -len(suffix_key)]
                rel_parent = str(p.parent.relative_to(root))
                results.append((p, rel_parent, stem, suffix_key))
                break
    return results


def compute_new_path(
    root: Path, rel_parent: str, stem: str, suffix_key: str,
) -> Path:
    """Compute the new path for a file."""
    modality, ext = _SUFFIX_MAP[suffix_key]
    if ext == ".npy":
        return root / "npy" / modality / rel_parent / f"{stem}{ext}"
    else:
        return root / modality / rel_parent / f"{stem}{ext}"


def migrate(root: Path, dry_run: bool = False) -> None:
    files = find_old_files(root)
    if not files:
        print(f"No old-layout files found in {root}")
        return

    print(f"Found {len(files)} files to migrate in {root}")

    moved = 0
    for old_path, rel_parent, stem, suffix_key in files:
        new_path = compute_new_path(root, rel_parent, stem, suffix_key)

        if dry_run:
            print(f"  {old_path.relative_to(root)} → {new_path.relative_to(root)}")
        else:
            new_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(str(old_path), str(new_path))
            moved += 1

    if dry_run:
        print(f"\nDry run: {len(files)} files would be moved.")
    else:
        print(f"Moved {moved} files.")

    # Clean up empty directories left behind.
    if not dry_run:
        _cleanup_empty_dirs(root)


def _cleanup_empty_dirs(root: Path) -> None:
    """Remove empty directories (bottom-up)."""
    removed = 0
    for dirpath, dirnames, filenames in os.walk(str(root), topdown=False):
        d = Path(dirpath)
        if d == root:
            continue
        # Skip new-layout dirs.
        try:
            rel = d.relative_to(root)
        except ValueError:
            continue
        if rel.parts and rel.parts[0] in _SKIP_DIRS:
            continue
        if not any(d.iterdir()):
            d.rmdir()
            removed += 1
    if removed:
        print(f"Cleaned up {removed} empty directories.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Migrate World-UAV-aug from prefix to directory layout",
    )
    parser.add_argument("root", type=Path, help="Path to World-UAV-aug directory")
    parser.add_argument("--dry-run", action="store_true",
                        help="Show what would be moved without doing it")
    args = parser.parse_args()

    if not args.root.is_dir():
        print(f"Error: {args.root} is not a directory")
        return

    migrate(args.root, dry_run=args.dry_run)


if __name__ == "__main__":
    main()