"""
检查 val.json 中每条记录的 image 和 mask 文件是否存在,
缺任一文件则删除该条,保存到新文件 (或原文件备份后覆盖)。

用法:
    python clean_val_json.py --json /path/to/val.json
    python clean_val_json.py --json /path/to/val.json --out /path/to/val_clean.json
    python clean_val_json.py --json /path/to/val.json --inplace   # 原地覆盖,自动备份为 .bak
"""
import argparse
import json
import os
from pathlib import Path


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--json", required=True, help="val.json 路径 (jsonl 格式)")
    ap.add_argument("--root", default=None,
                    help="image/mask 路径的根目录;默认 = val.json 所在目录")
    ap.add_argument("--out", default=None, help="输出路径,默认在原文件旁加 _clean")
    ap.add_argument("--inplace", action="store_true",
                    help="原地覆盖原文件 (会自动备份为 .bak)")
    args = ap.parse_args()

    json_path = Path(args.json)
    root = Path(args.root) if args.root else json_path.parent

    # 读 jsonl
    lines = [json.loads(l) for l in json_path.read_text().splitlines() if l.strip()]
    print(f"读取 {len(lines)} 条记录,根目录: {root}")

    kept, dropped = [], []
    for item in lines:
        img_rel = item["file_name"]
        # attn_list 是 [[label, mask_path], ...] 列表,所有 mask 都要存在
        mask_rels = [m[1] for m in item.get("attn_list", [])]

        img_ok = (root / img_rel).is_file()
        mask_ok = all((root / m).is_file() for m in mask_rels)

        if img_ok and mask_ok:
            kept.append(item)
        else:
            missing = []
            if not img_ok:
                missing.append(f"img:{img_rel}")
            for m in mask_rels:
                if not (root / m).is_file():
                    missing.append(f"mask:{m}")
            dropped.append((item.get("file_name"), missing))

    print(f"保留 {len(kept)} 条, 删除 {len(dropped)} 条")
    if dropped:
        print("\n删除示例 (最多前 5 条):")
        for fn, miss in dropped[:5]:
            print(f"  - {fn}  missing: {miss}")

    # 决定输出路径
    if args.inplace:
        backup = json_path.with_suffix(json_path.suffix + ".bak")
        if not backup.exists():            # 已有备份不覆盖,防止丢失原始版本
            os.rename(json_path, backup)
            print(f"\n原文件已备份: {backup}")
        else:
            print(f"\n[warn] 备份已存在,不再覆盖: {backup}")
        out_path = json_path
    else:
        out_path = Path(args.out) if args.out else \
            json_path.with_name(json_path.stem + "_clean" + json_path.suffix)

    with open(out_path, "w") as f:
        for item in kept:
            f.write(json.dumps(item) + "\n")
    print(f"已写出: {out_path}")


if __name__ == "__main__":
    main()


"""
python clean_val_json.py --json /home/jovyan/AURAD/detection/ex_validation/siim-ptx-test_prompt.json --root /home/jovyan/AURAD/detection/ex_validation/siim-ptx_test

"""