#!/usr/bin/env python3
"""Move existing query aux files (256x256) from World-UAV-aug to a backup folder.

Preserves the full directory structure so files can be restored if needed.
Uses os.rename (same filesystem) — no data copying, only inode renames.

Usage::

    # Dry run (default) — only counts, moves nothing:
    python -m src.scripts.move_query_aux

    # Actually move:
    python -m src.scripts.move_query_aux --execute

    # Restore (move back):
    python -m src.scripts.move_query_aux --restore --execute
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
from pathlib import Path

logger = logging.getLogger(__name__)

PARTITIONS = ("Country", "Terrain", "Rot")
AUX_SUFFIXES = ("_depth.png", "_edge.png", "_segm.png", "_chm.png")


def _is_query_aux(p: Path) -> bool:
    """Check if file is a query aux map (inside a query/ subtree)."""
    if p.suffix.lower() != ".png":
        return False
    if "query" not in p.parts:
        return False
    name = p.name
    return any(name.endswith(s) for s in AUX_SUFFIXES)


def move_query_aux(
    aux_root: Path,
    backup_root: Path,
    execute: bool = False,
    restore: bool = False,
) -> None:
    """Move query aux files between aux_root and backup_root."""

    if restore:
        src_root, dst_root = backup_root, aux_root
        action = "RESTORE"
    else:
        src_root, dst_root = aux_root, backup_root
        action = "BACKUP"

    if not src_root.exists():
        logger.error("Source root does not exist: %s", src_root)
        return

    logger.info("=== %s query aux files ===", action)
    logger.info("From: %s", src_root)
    logger.info("To:   %s", dst_root)
    logger.info("Mode: %s", "EXECUTE" if execute else "DRY RUN")

    moved = 0
    skipped = 0
    total_bytes = 0
    dirs_created: set[str] = set()

    for partition in PARTITIONS:
        pdir = src_root / partition
        if not pdir.exists():
            continue
        for p in sorted(pdir.rglob("*")):
            if not p.is_file():
                continue
            if not _is_query_aux(p):
                continue

            rel = p.relative_to(src_root)
            dst = dst_root / rel

            if dst.exists():
                skipped += 1
                continue

            total_bytes += p.stat().st_size
            moved += 1

            if execute:
                dst_dir = str(dst.parent)
                if dst_dir not in dirs_created:
                    dst.parent.mkdir(parents=True, exist_ok=True)
                    dirs_created.add(dst_dir)
                os.rename(p, dst)

            if moved % 100_000 == 0:
                logger.info("  ... %d files processed (%.1f GB)", moved, total_bytes / 1e9)

    logger.info("=== Done ===")
    logger.info("  Files %s: %d (%.1f GB)", "moved" if execute else "to move", moved, total_bytes / 1e9)
    if skipped:
        logger.info("  Skipped (already exists at destination): %d", skipped)
    if not execute:
        logger.info("  Run with --execute to actually move files.")


def main() -> None:
    parser = argparse.ArgumentParser(description="Move query aux files to/from backup")
    parser.add_argument("--aux-root", type=Path,
                        default=Path("/mnt/data1tb/cvgl_datasets/World-UAV-aug"))
    parser.add_argument("--backup-root", type=Path,
                        default=Path("/mnt/data1tb/cvgl_datasets/World-UAV-aug-query-256-backup"))
    parser.add_argument("--execute", action="store_true",
                        help="Actually move files (default is dry run)")
    parser.add_argument("--restore", action="store_true",
                        help="Restore files from backup back to aux_root")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)-7s %(message)s",
        datefmt="%H:%M:%S",
    )
    move_query_aux(args.aux_root, args.backup_root, args.execute, args.restore)


if __name__ == "__main__":
    main()
