from __future__ import annotations

import argparse
import json
from pathlib import Path

import yaml

from data.generate_rules import add_rule
from data.generate_scenes import generate_scene
from data.render_scene import render_scene


def _write_split(
    name: str,
    size: int,
    start: int,
    root: Path,
    seed: int,
    image_size: int,
    held_out: bool,
) -> None:
    path = root / f"{name}.jsonl"
    with path.open("w", encoding="utf-8") as handle:
        for scene_number in range(start, start + size):
            scene = generate_scene(
                scene_number, seed, held_out_composition=held_out
            )
            record = add_rule(scene, seed, held_out=held_out)
            relative_image = Path("images") / f"{record['scene_id']}.png"
            render_scene(record, root / relative_image, image_size=image_size)
            record["image"] = str(relative_image)
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")


def build(config_path: Path, small: bool = False, reuse_existing: bool = False) -> Path:
    config = yaml.safe_load(config_path.read_text(encoding="utf-8"))
    root = Path(config["data"]["root"])
    root.mkdir(parents=True, exist_ok=True)
    sizes = {
        "train": config["data"]["train_size"],
        "validation": config["data"]["validation_size"],
        "test_id": config["data"]["test_id_size"],
        "test_ood": config["data"]["test_ood_size"],
    }
    if small:
        sizes = {"train": 64, "validation": 16, "test_id": 16, "test_ood": 16}
    if reuse_existing and all(
        (root / f"{split}.jsonl").exists()
        and sum(1 for _ in (root / f"{split}.jsonl").open(encoding="utf-8")) == size
        for split, size in sizes.items()
    ):
        return root
    cursor = 0
    for split, size in sizes.items():
        _write_split(
            split,
            size,
            cursor,
            root,
            config["seed"],
            config["data"]["image_size"],
            held_out=split == "test_ood",
        )
        cursor += size
    return root


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml"))
    parser.add_argument("--small", action="store_true")
    parser.add_argument("--reuse-existing", action="store_true")
    args = parser.parse_args()
    print(build(args.config, small=args.small, reuse_existing=args.reuse_existing))


if __name__ == "__main__":
    main()