#!/usr/bin/env python3
"""
Inspect every <name>*_processed directory under a base for the self-match bug
(template-bank samples whose own match list contains themselves).

Output: a compact table + a JSON list of dirs needing fix at <out_json>.

Usage:
    python inspect_self_match_bug.py \
        --base /projects/p32685/RL_dataset \
        --out_json /projects/p32685/RL_dataset/_self_match_audit.json
"""
import argparse
import glob
import json
import os
import sys
from typing import List, Dict, Any


def list_processed_dirs(base: str) -> List[str]:
    """Every immediate child dir whose name ends in _processed (excluding obvious junk)."""
    out = []
    for name in sorted(os.listdir(base)):
        full = os.path.join(base, name)
        if not os.path.isdir(full):
            continue
        if not name.endswith("_processed"):
            continue
        # Skip symlinks: they alias another dir which we'll inspect directly.
        if os.path.islink(full):
            continue
        out.append(full)
    return out


def inspect_one(root: str) -> Dict[str, Any]:
    """Return audit fields for a single processed dir."""
    rep: Dict[str, Any] = {
        "root": root,
        "name": os.path.basename(root),
        "has_template_dir": False,
        "has_backup": False,
        "has_hf_dataset": False,
        "n_bank_npz": 0,
        "n_train_entries": 0,
        "n_val_entries": 0,
        "n_bank_in_train": 0,
        "n_bank_with_template_field": 0,
        "n_self_in_match_list": 0,        # any rank includes self
        "n_self_top1": 0,                 # top-1 match IS self (worst form)
        "n_train_missing_template": 0,    # any train row with empty/missing template
        "status": "UNKNOWN",
        "needs_fix": False,
        "errors": [],
    }
    tdir = os.path.join(root, "template")
    bdir = os.path.join(root, "backup_entries")
    hdir = os.path.join(root, "hf_dataset")
    rep["has_template_dir"] = os.path.isdir(tdir)
    rep["has_backup"] = os.path.isdir(bdir)
    rep["has_hf_dataset"] = os.path.isdir(hdir)

    if not (rep["has_template_dir"] and rep["has_backup"]):
        rep["status"] = "INCOMPLETE_DIR"
        return rep

    bank_files = glob.glob(os.path.join(tdir, "*.npz"))
    rep["n_bank_npz"] = len(bank_files)
    bank_stems = {os.path.splitext(os.path.basename(p))[0] for p in bank_files}

    # ---- inspect backup_entries (the source of truth for the template field) ----
    bp_train = os.path.join(bdir, "train_entries.json")
    bp_val = os.path.join(bdir, "val_entries.json")

    train_entries: List[Dict[str, Any]] = []
    val_entries: List[Dict[str, Any]] = []
    if os.path.isfile(bp_train):
        try:
            train_entries = json.load(open(bp_train))
        except Exception as e:
            rep["errors"].append(f"train_entries unreadable: {e}")
    if os.path.isfile(bp_val):
        try:
            val_entries = json.load(open(bp_val))
        except Exception as e:
            rep["errors"].append(f"val_entries unreadable: {e}")
    rep["n_train_entries"] = len(train_entries)
    rep["n_val_entries"] = len(val_entries)

    if not train_entries:
        rep["status"] = "NO_TRAIN_BACKUP"
        return rep

    for e in train_entries:
        stem = os.path.splitext(os.path.basename(e.get("image", "")))[0]
        is_bank = stem in bank_stems
        if is_bank:
            rep["n_bank_in_train"] += 1
        tpl = e.get("template")
        # Template field may be None, str (json), or list (after recover_from_backup re-encode)
        if isinstance(tpl, str):
            try:
                tpl = json.loads(tpl)
            except Exception:
                tpl = None
        if not isinstance(tpl, list) or len(tpl) == 0:
            rep["n_train_missing_template"] += 1
            continue
        if is_bank:
            rep["n_bank_with_template_field"] += 1
            ids = [t.get("template_id") for t in tpl]
            if stem in ids:
                rep["n_self_in_match_list"] += 1
            if ids and ids[0] == stem:
                rep["n_self_top1"] += 1

    # status logic
    if rep["n_bank_npz"] == 0:
        rep["status"] = "NO_BANK"
    elif rep["n_bank_in_train"] == 0:
        rep["status"] = "NO_BANK_OVERLAP"  # bank exists but no overlap with train (suspicious)
    elif rep["n_self_in_match_list"] > 0:
        rep["status"] = "BUGGY"
        rep["needs_fix"] = True
    elif rep["n_train_missing_template"] > 0:
        rep["status"] = "MISSING_TEMPLATE_FIELD"
        rep["needs_fix"] = True
    else:
        rep["status"] = "OK"

    return rep


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", required=True)
    ap.add_argument("--out_json", default=None)
    args = ap.parse_args()

    dirs = list_processed_dirs(args.base)
    print(f"Found {len(dirs)} *_processed directory under {args.base}")
    print()

    audits = []
    for d in dirs:
        rep = inspect_one(d)
        audits.append(rep)

    # ---- print summary table ----
    hdr = f"{'name':50s} {'bank':>5s} {'train':>6s} {'in_train':>8s} {'self∈list':>9s} {'self_top1':>9s} {'miss_tpl':>8s}  status"
    print(hdr)
    print("-" * len(hdr))
    for r in audits:
        print(
            f"{r['name']:50s} {r['n_bank_npz']:>5d} {r['n_train_entries']:>6d} "
            f"{r['n_bank_in_train']:>8d} {r['n_self_in_match_list']:>9d} "
            f"{r['n_self_top1']:>9d} {r['n_train_missing_template']:>8d}  {r['status']}"
        )

    buggy = [r for r in audits if r["needs_fix"]]
    print()
    print(f"NEEDS FIX: {len(buggy)} of {len(audits)}")
    for r in buggy:
        print(f"  - {r['name']:50s}  ({r['status']})")

    if args.out_json:
        with open(args.out_json, "w") as f:
            json.dump(audits, f, indent=2)
        print(f"\nFull audit saved to {args.out_json}")


if __name__ == "__main__":
    main()