#!/usr/bin/env python3.10

import argparse
import dataclasses
import datetime
import json
import logging
import re
from pathlib import Path


def parse_args():
    arger = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    arger.add_argument("dirs", help="Path to the directories", type=Path, nargs="+")
    arger.add_argument(
        "--flatten", help="Flatten directory structure", action="store_true"
    )
    return arger.parse_args()


def rename_children(source_dir: Path, flatten: bool) -> None:
    def natural_sort_key(path: Path) -> tuple[str, ...]:
        s = path.stem
        return tuple(
            f"{int(p):010d}" if p.isnumeric() else p
            for p in re.findall(r"(\D+|\d+)", s)
        )

    assert source_dir.is_dir()

    if flatten:
        files = [f for f in source_dir.rglob("*") if f.is_file()]
    else:
        files = [f for f in source_dir.iterdir() if f.is_file()]

    files = sorted(files, key=natural_sort_key)

    jobs = []
    for i, f in enumerate(files):
        ext = ".jpg" if f.suffix.lower() in {".jpeg", ".jpg"} else f.suffix
        target_dir = f.parent
        if flatten:
            target_dir = source_dir
        new_stem = source_dir.name
        target = target_dir / f"{new_stem}__{i:04d}{ext}"
        if target.is_file():
            logging.error(f"target already exists: {target}")
            continue
        jobs.append(MoveOp(source=f, target=target))
        f.rename(target)

    if not jobs:
        return

    write_undo_script(Path.cwd(), jobs)


@dataclasses.dataclass
class MoveOp:
    source: Path
    target: Path

    def to_dict(self) -> dict:
        return {
            "source": str(self.source),
            "target": str(self.target),
        }


def write_undo_script(save_dir: Path, ops: list[MoveOp]):
    now = datetime.datetime.now().isoformat().replace(":", "")
    save_path = save_dir / f"undo_{now}.json"
    with save_path.open("w", encoding="utf-8") as f:
        for it in ops:
            f.write(json.dumps(it.to_dict()) + "\n")


def main():
    args = parse_args()
    for d in args.dirs:
        rename_children(source_dir=d, flatten=args.flatten)


if __name__ == "__main__":
    main()