"""Convert the MuseumSCAT train.csv into HunyuanOCR conversation JSONL.

Reads train.csv + images/, writes:
  data/train.jsonl, data/val.jsonl   (HunyuanOCR raw-sample schema)
  data/images_resized/               (longest side capped, decoded once here
                                      instead of 8192px JPEGs every step)

The split is stratified over which fields are MISSING so the validation set
sees the same mix of empty and populated labels as training.

Usage:
  python prepare_data.py --dataset-dir /path/to/museumscat --out-dir ./data
"""

import argparse
import csv
import json
import os
import random
import sys

from PIL import Image

from common import MAX_SIDE, PROMPT, target_json


def build_conversation(image_path: str, date: str, locality: str) -> dict:
    return {
        "image_path": [image_path],
        "conversations": [
            {"from": "human", "value": "<image>\n" + PROMPT},
            {"from": "gpt", "value": target_json(date, locality)},
        ],
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset-dir", required=True, help="Directory containing train.csv and images/")
    parser.add_argument("--out-dir", default="data", help="Output directory")
    parser.add_argument("--val-fraction", type=float, default=0.1)
    parser.add_argument("--max-side", type=int, default=MAX_SIDE, help="Longest image side after resizing")
    parser.add_argument("--seed", type=int, default=13)
    args = parser.parse_args()

    train_csv = os.path.join(args.dataset_dir, "train.csv")
    images_dir = os.path.join(args.dataset_dir, "images")
    # Keyed on max_side so reruns with a different size regenerate images
    # instead of silently reusing the old resolution.
    resized_dir = os.path.join(args.out_dir, f"images_resized_{args.max_side}")
    os.makedirs(resized_dir, exist_ok=True)

    with open(train_csv, encoding="utf-8") as f:
        rows = list(csv.DictReader(f))

    samples = []
    for row in rows:
        image_file = row["image_file"].strip()
        date = row["verbatimDate"].strip()
        locality = row["verbatimLocality"].strip()
        src = os.path.join(images_dir, image_file)
        if not os.path.exists(src):
            print(f"skipping {image_file}: image not found", file=sys.stderr)
            continue
        if not date or not locality:
            print(f"skipping {image_file}: empty annotation", file=sys.stderr)
            continue
        samples.append({"image_file": image_file, "date": date, "locality": locality})

    # Resize once, up front.
    for sample in samples:
        src = os.path.join(images_dir, sample["image_file"])
        dst = os.path.join(resized_dir, sample["image_file"])
        sample["resized_path"] = os.path.abspath(dst)
        if os.path.exists(dst):
            continue
        with Image.open(src) as im:
            im = im.convert("RGB")
            im.thumbnail((args.max_side, args.max_side), Image.LANCZOS)
            im.save(dst, quality=92)

    # Stratify by (date missing, locality missing).
    strata = {}
    for sample in samples:
        key = (sample["date"] == "MISSING", sample["locality"] == "MISSING")
        strata.setdefault(key, []).append(sample)

    rng = random.Random(args.seed)
    train_set, val_set = [], []
    for key in sorted(strata):
        group = strata[key]
        rng.shuffle(group)
        n_val = max(1, round(len(group) * args.val_fraction))
        val_set.extend(group[:n_val])
        train_set.extend(group[n_val:])
    rng.shuffle(train_set)

    for name, subset in (("train", train_set), ("val", val_set)):
        path = os.path.join(args.out_dir, f"{name}.jsonl")
        with open(path, "w", encoding="utf-8") as f:
            for sample in subset:
                record = build_conversation(sample["resized_path"], sample["date"], sample["locality"])
                record["image_file"] = sample["image_file"]
                f.write(json.dumps(record, ensure_ascii=False) + "\n")
        print(f"{path}: {len(subset)} samples")

    n_missing_date = sum(1 for s in val_set if s["date"] == "MISSING")
    n_missing_loc = sum(1 for s in val_set if s["locality"] == "MISSING")
    print(f"val MISSING dates: {n_missing_date}, MISSING localities: {n_missing_loc}")


if __name__ == "__main__":
    main()