#!/usr/bin/env python3
"""
Path A in-place repair of one already-built <name>*_processed dir.

What it does (and ONLY this — bank/.npz files untouched):
  1. Load the existing template bank from <root>/template/*.npz.
  2. Re-run patched precompute_for_entries (self-exclusion ON, --overwrite ON) on
     backup_entries/{train,val}_entries.json, regenerating the `template` field
     for every SFT sample.
  3. recover_from_backup() rebuilds <root>/hf_dataset/ from the updated backups.

After this, every bank sample's match list is guaranteed not to contain itself.
SFT samples that were missing a template entirely also get one.

Usage:
    python repair_self_match.py --root /projects/p32685/RL_dataset/ACDC_10percent_processed
"""
import argparse
import json
import os
import sys
from pathlib import Path
from types import SimpleNamespace

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 packing_HF_dataset import GenHFDataset
from template_manager import TemplateManager
from precompute_template import precompute_for_entries


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", required=True, type=str,
                    help="Existing <name>*_processed dir (must contain template/, backup_entries/)")
    ap.add_argument("--device", type=str, default="cuda:0")
    # NOTE: top-1 only (was previously top-k for downstream random.choice; that mechanism
    # is no longer in use, so we just store the single best match).
    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 repair_one(root, device="cuda:0", topk=1, match_strategy="partial", max_search=20):
    root = os.path.abspath(root)
    template_dir = os.path.join(root, "template")
    bdir = os.path.join(root, "backup_entries")
    bp_train = os.path.join(bdir, "train_entries.json")
    bp_val = os.path.join(bdir, "val_entries.json")

    if not os.path.isdir(template_dir):
        raise FileNotFoundError(f"missing template/ under {root}")
    if not os.path.isfile(bp_train):
        raise FileNotFoundError(f"missing backup_entries/train_entries.json under {root}")

    print(f"\n[REPAIR] {root}")

    # ---- step 1: load existing template bank ----
    tm = TemplateManager(device=device, distance_metric="cosine")
    templates, faiss_index, template_ids_list = tm._load_templates(template_dir)
    print(f"  loaded {len(templates)} templates from {template_dir}")

    # ---- step 2: re-run precompute_for_entries with --overwrite=True ----
    train_entries = json.load(open(bp_train))
    val_entries = json.load(open(bp_val)) if os.path.isfile(bp_val) else []

    # NOTE: a Python class body has its own scope and cannot see the function locals,
    # so `topk = topk` inside `class _Args:` raised NameError. Use SimpleNamespace instead.
    cfg = SimpleNamespace(
        overwrite=True,                # always re-write the template field
        topk=topk,                     # default 1 — single best match only
        match_strategy=match_strategy,
        max_search=max_search,
    )

    u_train = precompute_for_entries(
        "train", train_entries, templates, faiss_index, template_ids_list, tm, cfg,
    )
    print(f"  precompute updated train rows: {u_train}/{len(train_entries)}")

    u_val = precompute_for_entries(
        "val", val_entries, templates, faiss_index, template_ids_list, tm, cfg,
    )
    print(f"  precompute updated val rows  : {u_val}/{len(val_entries)}")

    with open(bp_train, "w") as f:
        json.dump(train_entries, f, indent=2)
    if os.path.isfile(bp_val):
        with open(bp_val, "w") as f:
            json.dump(val_entries, f, indent=2)
    print(f"  saved updated backup_entries")

    # ---- step 3: rebuild hf_dataset from updated backups ----
    gen = GenHFDataset()
    print(f"  recover_from_backup -> rebuilding hf_dataset")
    gen.recover_from_backup(output_dir=root)

    # ---- step 4: post-fix verification ----
    import glob
    bank_stems = {os.path.splitext(os.path.basename(p))[0] for p in glob.glob(os.path.join(template_dir, "*.npz"))}
    train_after = json.load(open(bp_train))
    self_top1, self_in_list, miss = 0, 0, 0
    for e in train_after:
        stem = os.path.splitext(os.path.basename(e["image"]))[0]
        if stem not in bank_stems:
            continue
        tpl = e.get("template")
        if not isinstance(tpl, list) or not tpl:
            miss += 1
            continue
        ids = [t.get("template_id") for t in tpl]
        if stem in ids:
            self_in_list += 1
        if ids[0] == stem:
            self_top1 += 1
    print(f"  POST-FIX: bank={len(bank_stems)}  self_in_list={self_in_list}  self_top1={self_top1}  missing_template={miss}")
    if self_in_list == 0 and self_top1 == 0:
        print(f"  ✓ {os.path.basename(root)}: CLEAN")
    else:
        print(f"  ✗ {os.path.basename(root)}: STILL HAS VIOLATIONS")
    return self_in_list == 0 and self_top1 == 0


def main():
    args = parse_args()
    repair_one(args.root, device=args.device, topk=args.topk,
               match_strategy=args.match_strategy, max_search=args.max_search)


if __name__ == "__main__":
    main()