#!/usr/bin/env python3

from pathlib import Path
import json
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("bids_root")
parser.add_argument("subjects", nargs="+")
parser.add_argument("--session", default="001")
args = parser.parse_args()

root = Path(args.bids_root)
ses = f"ses-{args.session}"

for sub in args.subjects:
    ses_dir = root / sub / ses
    fmap_dir = ses_dir / "fmap"
    dwi_dir = ses_dir / "dwi"

    pairs = [
        ("b1200", f"{sub}_{ses}_acq-b1200_dir-PA_dwi.nii.gz"),
        ("mshell", f"{sub}_{ses}_acq-mshell_dir-PA_dwi.nii.gz"),
    ]

    for acq, dwi_name in pairs:
        dwi_file = dwi_dir / dwi_name
        epi_json = fmap_dir / f"{sub}_{ses}_acq-{acq}_dir-AP_epi.json"

        if not dwi_file.exists() or not epi_json.exists():
            print(f"SKIP {sub} {acq}: missing DWI or fmap JSON")
            continue

        with epi_json.open("r", encoding="utf-8") as f:
            data = json.load(f)

        data["IntendedFor"] = [f"{ses}/dwi/{dwi_name}"]

        with epi_json.open("w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)
            f.write("\n")

        print(f"UPDATED {epi_json}")