"""
Batch migrate omni.isaac.* imports to use core.compat module.

Usage:
    python migrate/migrate_imports.py --dry-run   # preview changes
    python migrate/migrate_imports.py              # apply changes
"""
import argparse
import os
import re

# Files to process (relative to project root)
SCAN_DIRS = [
    "workflows/simbox/core/cameras",
    "workflows/simbox/core/controllers",
    "workflows/simbox/core/objects",
    "workflows/simbox/core/robots",
    "workflows/simbox/core/skills",
    "workflows/simbox/core/tasks",
    "workflows/simbox/core/utils",
    "workflows/simbox/solver",
    "workflows/",  # simbox_dual_workflow.py
    "nimbus_extension/components/load",
]

# Skip these files/dirs
SKIP = ["curobo", "__pycache__", "compat.py", "tools/art"]

# Mapping: old import pattern -> compat import
# Format: (regex_pattern, replacement_imports_from_compat)
IMPORT_MAP = {
    # Full module imports
    r"from omni\.isaac\.core import World": "from core.compat import World",
    r"from omni\.isaac\.core\.controllers import BaseController": "from core.compat import BaseController",
    r"from omni\.isaac\.core\.tasks import BaseTask": "from core.compat import BaseTask",
    r"from omni\.isaac\.core\.robots\.robot import Robot": "from core.compat import Robot",
    r"from omni\.isaac\.core\.robots import Robot": "from core.compat import Robot",
    r"from omni\.isaac\.core\.prims import XFormPrim": "from core.compat import XFormPrim",
    r"from omni\.isaac\.core\.prims import RigidPrim": "from core.compat import RigidPrim",
    r"from omni\.isaac\.core\.prims import GeometryPrim": "from core.compat import GeometryPrim",
    r"from omni\.isaac\.core\.prims import RigidContactView": "from core.compat import RigidContactView",
    r"from omni\.isaac\.core\.prims import RigidContactView, XFormPrim": "from core.compat import RigidContactView, XFormPrim",
    r"from omni\.isaac\.core\.articulations\.articulation import Articulation": "from core.compat import Articulation",
    r"from omni\.isaac\.core\.objects\.cylinder import VisualCylinder": "from core.compat import VisualCylinder",
    r"from omni\.isaac\.core\.objects import cuboid, sphere": "from core.compat import cuboid, sphere",
    r"from omni\.isaac\.core\.objects import cuboid": "from core.compat import cuboid",
    r"from omni\.isaac\.core\.objects import sphere": "from core.compat import sphere",
    r"from omni\.isaac\.core\.materials import PreviewSurface": "from core.compat import PreviewSurface",
    r"from omni\.isaac\.core\.materials import OmniPBR": "from core.compat import OmniPBR",
    r"from omni\.isaac\.core\.materials\.omni_pbr import OmniPBR.*": "from core.compat import OmniPBR",
    r"from omni\.isaac\.core\.materials import OmniGlass, OmniPBR": "from core.compat import OmniGlass, OmniPBR",
    r"from omni\.isaac\.core\.materials import OmniGlass": "from core.compat import OmniGlass",
    r"from omni\.isaac\.core\.scenes\.scene import Scene": "from core.compat import Scene",
    r"from omni\.isaac\.sensor import Camera": "from core.compat import Camera",

    # Utils
    r"from omni\.isaac\.core\.utils\.prims import get_prim_at_path\b(?!,)": "from core.compat import get_prim_at_path",
    r"from omni\.isaac\.core\.utils\.prims import create_prim, get_prim_at_path": "from core.compat import create_prim, get_prim_at_path",
    r"from omni\.isaac\.core\.utils\.prims import create_prim, is_prim_path_valid": "from core.compat import create_prim, is_prim_path_valid",
    r"from omni\.isaac\.core\.utils\.prims import is_prim_path_valid": "from core.compat import is_prim_path_valid",
    r"from omni\.isaac\.core\.utils\.prims import create_prim\b": "from core.compat import create_prim",
    r"from omni\.isaac\.core\.utils\.transformations import get_relative_transform\b(?!,)": "from core.compat import get_relative_transform",
    r"from omni\.isaac\.core\.utils\.transformations import pose_from_tf_matrix": "from core.compat import pose_from_tf_matrix",
    r"from omni\.isaac\.core\.utils\.stage import get_current_stage": "from core.compat import get_current_stage",
    r"from omni\.isaac\.core\.utils\.stage import add_reference_to_stage": "from core.compat import add_reference_to_stage",
    r"from omni\.isaac\.core\.utils\.stage import get_stage_units": "from core.compat import get_stage_units",
    r"from omni\.isaac\.core\.utils\.xforms import get_world_pose": "from core.compat import get_world_pose",
    r"from omni\.isaac\.core\.utils\.types import ArticulationAction": "from core.compat import ArticulationAction",
    r"from omni\.isaac\.core\.utils\.viewports import set_camera_view": "from core.compat import set_camera_view",
    r"from omni\.isaac\.core\.utils\.semantics import add_update_semantics": "from core.compat import add_update_semantics",
    r"from omni\.isaac\.core\.utils\.string import find_unique_string_name": "from core.compat import find_unique_string_name",
    r"from omni\.isaac\.core\.utils\.extensions import enable_extension": "from core.compat import enable_extension",
    r"from omni\.isaac\.core\.utils\.nucleus import get_assets_root_path as nucleus_path": "from core.compat import nucleus_path",
}

# Multi-line import blocks that need special handling
MULTILINE_PATTERNS = [
    # get_relative_transform, pose_from_tf_matrix block
    (
        r"from omni\.isaac\.core\.utils\.transformations import \(\s*\n\s*get_relative_transform,\s*\n\s*pose_from_tf_matrix,?\s*\n\s*\)",
        "from core.compat import get_relative_transform, pose_from_tf_matrix",
    ),
    # prims multi-import blocks
    (
        r"from omni\.isaac\.core\.utils\.prims import \(\s*\n\s*get_prim_at_path,\s*\n\s*create_prim,\s*\n\s*is_prim_path_valid,?\s*\n\s*\)",
        "from core.compat import get_prim_at_path, create_prim, is_prim_path_valid",
    ),
    (
        r"from omni\.isaac\.core\.utils\.prims import \(\s*\n[^)]+\)",
        None,  # skip complex blocks, handle manually
    ),
]


def should_skip(filepath):
    for s in SKIP:
        if s in filepath:
            return True
    return False


def migrate_file(filepath, dry_run=False):
    with open(filepath, "r") as f:
        content = f.read()

    original = content
    changes = []

    # Handle multi-line imports first
    for pattern, replacement in MULTILINE_PATTERNS:
        if replacement is None:
            continue
        matches = list(re.finditer(pattern, content))
        for m in matches:
            changes.append(f"  MULTI: {m.group()[:60]}... -> {replacement}")
        content = re.sub(pattern, replacement, content)

    # Handle single-line imports
    for pattern, replacement in IMPORT_MAP.items():
        matches = list(re.finditer(pattern, content))
        for m in matches:
            changes.append(f"  {m.group().strip()} -> {replacement}")
        content = re.sub(pattern, replacement, content)

    if content != original:
        if not dry_run:
            with open(filepath, "w") as f:
                f.write(content)
        return changes
    return []


def main():
    parser = argparse.ArgumentParser(description="Migrate omni.isaac imports to core.compat")
    parser.add_argument("--dry-run", action="store_true", help="Preview changes without modifying files")
    parser.add_argument("--root", default=".", help="Project root directory")
    args = parser.parse_args()

    total_files = 0
    total_changes = 0

    for scan_dir in SCAN_DIRS:
        full_dir = os.path.join(args.root, scan_dir)
        if not os.path.exists(full_dir):
            continue

        if os.path.isfile(full_dir):
            files = [full_dir]
        else:
            files = []
            for root, dirs, filenames in os.walk(full_dir):
                for fn in filenames:
                    if fn.endswith(".py"):
                        files.append(os.path.join(root, fn))

        for filepath in sorted(files):
            if should_skip(filepath):
                continue

            rel = os.path.relpath(filepath, args.root)
            changes = migrate_file(filepath, args.dry_run)
            if changes:
                total_files += 1
                total_changes += len(changes)
                print(f"\n{rel}:")
                for c in changes:
                    print(c)

    action = "would change" if args.dry_run else "changed"
    print(f"\n{'='*60}")
    print(f"Done: {total_changes} imports {action} in {total_files} files.")
    if args.dry_run:
        print("Run without --dry-run to apply changes.")


if __name__ == "__main__":
    main()