#!/usr/bin/env python3
"""
For one <name>_*_processed dir, rewrite broken image / label_mask paths in
backup_entries to match the current disk layout, THEN run the in-place
self-match repair (precompute_for_entries with self-exclusion + recover_from_backup).

Two rewrite rules, applied in order:
  1. /projects/p32591/  ->  /projects/p32685/        (project-allocation move)
  2. RL_dataset/MSD_taskXX/TaskXX_<Name>/  ->  RL_dataset/MSD_taskXX_<Name>/
     (older path layout had a nested MSD_taskXX/TaskXX_<Name> dir; current disk has
     a flat MSD_taskXX_<Name> dir)

Verifies the rewritten path actually exists; if any sample's path fails to
resolve, the entry is flagged in a side-log so the user can fix it manually.

Usage:
    python fix_paths_and_repair.py --root /projects/p32685/RL_dataset/MSD_task02_Heart_10percent_processed
"""
import argparse
import json
import os
import re
import sys
from pathlib import Path

THIS_DIR = Path(__file__).resolve().parent
VESSA_ROOT = THIS_DIR.parent
for p in [str(THIS_DIR),
          str(VESSA_ROOT / "train_vessa"),
          str(VESSA_ROOT / "train_vessa" / "tools")]:
    if p not in sys.path:
        sys.path.insert(0, p)

from repair_self_match import repair_one


# Pattern 2 — collapse MSD_taskXX/TaskXX_<Name>/ -> MSD_taskXX_<Name>/
_MSD_RE = re.compile(r"/RL_dataset/(MSD_task\d+)/(Task\d+_[^/]+)/")


def rewrite_path(p: str) -> str:
    if not isinstance(p, str):
        return p
    # Rule 1: p32591 -> p32685
    p2 = p.replace("/projects/p32591/", "/projects/p32685/")
    # Rule 2: MSD_taskXX/TaskXX_<Name>/ -> MSD_taskXX_<Name>/
    m = _MSD_RE.search(p2)
    if m:
        # Use the second group (Task02_Heart) and turn it into MSD_task02_Heart
        # via the first group's task number.
        msd_prefix = m.group(1)            # e.g. MSD_task02
        task_with_name = m.group(2)        # e.g. Task02_Heart
        # Strip "Task02_" and prepend "MSD_task02_" -> "MSD_task02_Heart"
        try:
            _, name = task_with_name.split("_", 1)
        except ValueError:
            return p2
        new_dir = f"{msd_prefix}_{name}"
        p2 = _MSD_RE.sub(f"/RL_dataset/{new_dir}/", p2, count=1)
    return p2


def fix_paths_in_backup(root: str) -> dict:
    bdir = os.path.join(root, "backup_entries")
    stats = {"train_total": 0, "train_changed": 0, "train_unresolved": 0,
             "val_total": 0, "val_changed": 0, "val_unresolved": 0,
             "unresolved_examples": []}
    for split in ("train", "val"):
        p = os.path.join(bdir, f"{split}_entries.json")
        if not os.path.isfile(p):
            continue
        entries = json.load(open(p))
        changed = 0
        unresolved = 0
        for e in entries:
            for key in ("image", "label_mask"):
                old = e.get(key)
                new = rewrite_path(old)
                if new != old:
                    e[key] = new
                    changed += 1
                if not os.path.exists(new):
                    unresolved += 1
                    if len(stats["unresolved_examples"]) < 5:
                        stats["unresolved_examples"].append(new)
        with open(p, "w") as f:
            json.dump(entries, f, indent=2)
        stats[f"{split}_total"] = len(entries) * 2  # 2 paths per entry
        stats[f"{split}_changed"] = changed
        stats[f"{split}_unresolved"] = unresolved
    return stats


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", required=True, type=str)
    ap.add_argument("--device", type=str, default="cuda:0")
    ap.add_argument("--topk", type=int, default=1)
    ap.add_argument("--match_strategy", type=str, default="partial",
                    choices=["partial", "complete"])
    ap.add_argument("--max_search", type=int, default=20)
    return ap.parse_args()


def main():
    args = parse_args()
    root = os.path.abspath(args.root)
    print(f"\n[FIX-AND-REPAIR] {root}")

    # Step 1 — rewrite broken paths in backup_entries.
    stats = fix_paths_in_backup(root)
    print(f"  paths changed   : train={stats['train_changed']:5d}  val={stats['val_changed']:5d}")
    print(f"  paths unresolved: train={stats['train_unresolved']:5d}  val={stats['val_unresolved']:5d}")
    if stats["unresolved_examples"]:
        print(f"  example unresolved paths (first 5):")
        for p in stats["unresolved_examples"]:
            print(f"    {p}")

    # Step 2 — re-run repair on the (now hopefully resolvable) backup_entries.
    repair_one(root, device=args.device, topk=args.topk,
               match_strategy=args.match_strategy, max_search=args.max_search)


if __name__ == "__main__":
    main()