#!/usr/bin/env python3
"""
ROM library dedup: keeps the best regional variant of each game
(USA > World > Europe > Japan), deletes the rest.

Usage:
  python3 rom-dedup.py                   # dry run (no changes)
  python3 rom-dedup.py --execute         # actually delete duplicates
  python3 rom-dedup.py --include-arcade  # also process the arcade dir
  python3 rom-dedup.py --console snes    # only process one console dir
"""

import os
import re
import sys
import argparse
from collections import defaultdict
from pathlib import Path

# ---- config ----
LIBRARY_PATH = "/path/to/roms"  # one subdirectory per console, zip files inside
LOG_PATH = "dedup_log.txt"      # run log destination
# ----------------

# tags marking non-standard releases; these lose to normal releases
PENALTY_TAGS = [
    "virtual console", "retro-bit", "retro-bit generations",
    "disney classic games", "final cut", "pirate", "digital",
    "ndsi enhanced", "gamecube edition", "arcade", "unl",
    "aftermarket", "homebrew", "proto", "beta", "demo", "sample",
    "program", "promo",
]


def score_file(filename: str) -> int:
    """Higher score = preferred (kept)."""
    name = filename.lower()

    if "(usa)" in name:
        score = 100
    elif "(world)" in name:
        score = 80
    elif "(usa," in name or "(en," in name:
        score = 65
    elif "(europe)" in name:
        score = 50
    elif "(australia)" in name:
        score = 45
    elif "(uk)" in name or "(united kingdom)" in name:
        score = 44
    elif "(japan)" in name or "(ja)" in name:
        score = 20
    elif "(china)" in name or "(zh)" in name:
        score = 10
    elif "(korea)" in name or "(ko)" in name:
        score = 10
    else:
        score = 5

    # higher revision = errata fixed, small bonus within the region bucket
    rev_match = re.search(r'\(rev\s*([0-9]+)\)', name)
    if rev_match:
        score += int(rev_match.group(1))

    for tag in PENALTY_TAGS:
        if tag in name:
            score -= 30
            break

    # small penalty per parenthetical tag group: prefer cleaner releases
    score -= len(re.findall(r'\([^)]+\)', filename))

    return score


def get_base_name(filename: str) -> str:
    """Strip extension and trailing (tag) groups to get the game's base name."""
    name = Path(filename).stem
    name = re.sub(r'(\s*\([^)]*\))+$', '', name).strip()
    return name.lower()


def find_duplicates(console_path: str) -> dict:
    """base_name -> list of filenames sorted best-first; only groups with >1 file."""
    try:
        files = [f for f in os.listdir(console_path) if f.lower().endswith('.zip')]
    except PermissionError:
        return {}

    groups = defaultdict(list)
    for f in files:
        groups[get_base_name(f)].append(f)

    duplicates = {}
    for base, file_list in groups.items():
        if len(file_list) > 1:
            duplicates[base] = sorted(file_list, key=score_file, reverse=True)

    return duplicates


def human_size(total_bytes: int) -> str:
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if total_bytes < 1024:
            return f"{total_bytes:.1f} {unit}"
        total_bytes /= 1024
    return f"{total_bytes:.1f} PB"


def main():
    parser = argparse.ArgumentParser(description="Deduplicate ROM library")
    parser.add_argument('--execute', action='store_true',
                        help="Actually delete files (default: dry run)")
    parser.add_argument('--include-arcade', action='store_true',
                        help="Include arcade ROMs (risky: MAME sets are not true duplicates)")
    parser.add_argument('--console', type=str, default=None,
                        help="Only process a specific console directory name")
    args = parser.parse_args()

    dry_run = not args.execute
    mode = "DRY RUN" if dry_run else "EXECUTE"

    library = Path(LIBRARY_PATH)
    if not library.exists():
        print(f"ERROR: Library path not found: {LIBRARY_PATH}")
        sys.exit(1)

    if args.console:
        consoles = [library / args.console]
        if not consoles[0].exists():
            print(f"ERROR: Console directory not found: {consoles[0]}")
            sys.exit(1)
    else:
        consoles = sorted([d for d in library.iterdir() if d.is_dir()])

    lines = []
    total_delete_count = 0
    total_delete_bytes = 0

    header = f"=== ROM Dedup [{mode}] | {LIBRARY_PATH} ==="
    print(header)
    lines.append(header)

    for console_dir in consoles:
        console_name = console_dir.name

        if console_name == "arcade" and not args.include_arcade:
            msg = f"\n[{console_name}] SKIPPED (arcade: use --include-arcade to process)"
            print(msg)
            lines.append(msg)
            continue

        dupes = find_duplicates(str(console_dir))
        if not dupes:
            continue

        console_header = f"\n[{console_name}] {len(dupes)} game(s) with duplicates:"
        print(console_header)
        lines.append(console_header)

        for base, scored_files in sorted(dupes.items()):
            keep = scored_files[0]
            to_delete = scored_files[1:]

            keep_msg = f"  KEEP   [{score_file(keep):4d}] {keep}"
            print(keep_msg)
            lines.append(keep_msg)

            for f in to_delete:
                fpath = console_dir / f
                try:
                    size = fpath.stat().st_size
                except OSError:
                    size = 0

                delete_msg = f"  DELETE [{score_file(f):4d}] {f}  ({human_size(size)})"
                print(delete_msg)
                lines.append(delete_msg)

                total_delete_count += 1
                total_delete_bytes += size

                if not dry_run:
                    try:
                        fpath.unlink()
                    except OSError as e:
                        err = f"  ERROR deleting {fpath}: {e}"
                        print(err)
                        lines.append(err)

    summary = (
        f"\n{'='*60}\n"
        f"{'[DRY RUN] Would delete' if dry_run else 'Deleted'}: "
        f"{total_delete_count} files  ({human_size(total_delete_bytes)} freed)\n"
        f"{'='*60}"
    )
    print(summary)
    lines.append(summary)

    if dry_run:
        tip = "\nRe-run with --execute to apply deletions."
        print(tip)
        lines.append(tip)

    try:
        with open(LOG_PATH, 'w') as f:
            f.write('\n'.join(lines))
        print(f"\nLog written to: {LOG_PATH}")
    except OSError as e:
        print(f"\nWarning: could not write log: {e}")


if __name__ == "__main__":
    main()

# src: ezra, contact: ezkru69@mail.com