"""Concatenate gallery JSONLs into one training/eval file in a new directory, rewriting image paths so they
resolve relative to the output file (train.py / evaluate.py take a single --data file).

Each input is `path` (all records) or `path@N` (N records sampled with --seed, input order kept).
Every output record gets meta.mix_source = the input path (relative to the output dir).

usage: python scripts/make_mix.py out.jsonl in1.jsonl in2.jsonl@300 [--seed 0]
"""
import json, os, random, sys
from collections import Counter


def main(argv):
    seed = 0
    if "--seed" in argv:
        i = argv.index("--seed"); seed = int(argv[i + 1]); del argv[i:i + 2]
    dst, inputs = argv[0], argv[1:]
    odir = os.path.dirname(os.path.abspath(dst)); os.makedirs(odir, exist_ok=True)
    rng = random.Random(seed); n_out = 0; per = Counter(); seen = set()
    with open(dst, "w", encoding="utf-8", newline="\n") as out:
        for spec in inputs:
            path, n = (spec.rsplit("@", 1) + [None])[:2] if "@" in spec else (spec, None)
            with open(path, encoding="utf-8") as f:
                rows = [json.loads(l) for l in f if l.strip()]
            if n is not None:
                keep = set(rng.sample(range(len(rows)), min(int(n), len(rows))))
                rows = [r for i, r in enumerate(rows) if i in keep]
            idir = os.path.dirname(os.path.abspath(path))
            src_rel = os.path.relpath(os.path.abspath(path), odir).replace("\\", "/")
            for r in rows:
                if r["id"] in seen:
                    sys.exit(f"duplicate id {r['id']} ({spec})")
                seen.add(r["id"])
                r["images"] = [im if os.path.isabs(im) else os.path.relpath(os.path.join(idir, im), odir).replace("\\", "/")
                               for im in r["images"]]
                missing = [im for im in r["images"] if not os.path.exists(os.path.join(odir, im))]
                if missing:
                    sys.exit(f"missing image {missing} ({spec})")
                r.setdefault("meta", {})["mix_source"] = src_rel
                out.write(json.dumps(r, ensure_ascii=False) + "\n"); n_out += 1; per[src_rel] += 1
    n_box = 0
    with open(dst, encoding="utf-8") as f:
        for l in f:
            r = json.loads(l); n_box += any(q.get("type") == "box" for q in r["questions"].values())
    print(json.dumps({"out": dst, "records": n_out, "records_with_box_questions": n_box, "per_input": per, "seed": seed}))


if __name__ == "__main__":
    main(sys.argv[1:])