#!/usr/bin/env python3
"""
Rewrite absolute paths embedded in processed datasets (e.g. after copying project p32591 -> p32685).

Updates:
  - backup_entries/*.json (image, label_mask string fields; template JSON if present)
  - template/template_summary.json
  - hf_dataset/ (HuggingFace DatasetDict on disk) columns image, label_mask
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
import tempfile


def replace_in_obj(obj, old: str, new: str):
    if isinstance(obj, str):
        return obj.replace(old, new)
    if isinstance(obj, list):
        return [replace_in_obj(x, old, new) for x in obj]
    if isinstance(obj, dict):
        return {k: replace_in_obj(v, old, new) for k, v in obj.items()}
    return obj


def fix_backup_entries(processed_dir: str, old: str, new: str) -> int:
    n = 0
    bdir = os.path.join(processed_dir, "backup_entries")
    for name in ("train_entries.json", "val_entries.json"):
        path = os.path.join(bdir, name)
        if not os.path.isfile(path):
            continue
        with open(path) as f:
            data = json.load(f)
        data = replace_in_obj(data, old, new)
        with open(path, "w") as f:
            json.dump(data, f, indent=2)
        n += 1
    return n


def fix_template_summary(processed_dir: str, old: str, new: str) -> bool:
    path = os.path.join(processed_dir, "template", "template_summary.json")
    if not os.path.isfile(path):
        return False
    with open(path) as f:
        data = json.load(f)
    data = replace_in_obj(data, old, new)
    with open(path, "w") as f:
        json.dump(data, f, indent=2)
    return True


def fix_hf_dataset(processed_dir: str, old: str, new: str) -> bool:
    hf_path = os.path.join(processed_dir, "hf_dataset")
    if not os.path.isdir(hf_path):
        return False
    from datasets import DatasetDict, load_from_disk

    ds = load_from_disk(hf_path)

    def fix_row(ex):
        out = dict(ex)
        for key in ("image", "label_mask"):
            if key in out and isinstance(out[key], str):
                out[key] = out[key].replace(old, new)
        return out

    new_splits = {}
    for name, split in ds.items():
        new_splits[name] = split.map(fix_row)

    parent = os.path.dirname(hf_path.rstrip(os.sep))
    tmp = tempfile.mkdtemp(prefix=".hf_dataset_rewrite_", dir=parent)
    try:
        DatasetDict(new_splits).save_to_disk(tmp)
        shutil.rmtree(hf_path)
        shutil.move(tmp, hf_path)
    except Exception:
        shutil.rmtree(tmp, ignore_errors=True)
        raise
    return True


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument(
        "--processed_glob",
        default="/projects/p32685/RL_dataset/lungx_*percent_processed",
        help="Glob of processed dataset dirs",
    )
    ap.add_argument(
        "--old_prefix",
        default="/projects/p32591/RL_dataset",
        help="Path prefix to replace",
    )
    ap.add_argument(
        "--new_prefix",
        default="/projects/p32685/RL_dataset",
        help="Replacement prefix",
    )
    args = ap.parse_args()

    import glob

    dirs = sorted(glob.glob(args.processed_glob))
    if not dirs:
        print(f"No dirs matched: {args.processed_glob}", file=sys.stderr)
        return 1

    old, new = args.old_prefix, args.new_prefix
    for d in dirs:
        print(d)
        fix_backup_entries(d, old, new)
        print("  backup_entries: ok")
        if fix_template_summary(d, old, new):
            print("  template_summary.json: ok")
        else:
            print("  template_summary.json: skip")
        if fix_hf_dataset(d, old, new):
            print("  hf_dataset: ok")
        else:
            print("  hf_dataset: skip")

    print("Done. Re-run Step 3 if needed.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())