#!/usr/bin/env python3
from __future__ import annotations

import json
import random
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
OUTPUT = ROOT / "parser" / "output"
HERE = Path(__file__).resolve().parent
DATASET = Path("/Users/jerry/linkhand/cad/SLDPRT-NET/sub_datset/sub_datset")
TARGET_COUNT = 50


def main() -> None:
    rows: list[dict] = []
    for path in sorted(OUTPUT.glob("*.json")):
        payload = json.loads(path.read_text())
        stem = Path(payload["provenance"]["source_path"]).stem
        image = DATASET / "image" / f"{stem}.png"
        description = DATASET / "Des_Text" / f"{stem}.txt"
        reference = DATASET / "model_step" / f"{stem}.step"
        if image.is_file() and description.is_file() and reference.is_file():
            rows.append(
                {
                    "id": stem,
                    "family": payload["design_ir"]["part_family"],
                    "case_json": str(path.resolve()),
                    "image": str(image),
                    "description": str(description),
                    "reference_step": str(reference),
                }
            )

    by_family: dict[str, list[dict]] = defaultdict(list)
    for row in rows:
        by_family[row["family"]].append(row)
    rng = random.Random(20260723)
    selected: list[dict] = []
    total = len(rows)
    for family, family_rows in sorted(by_family.items()):
        rng.shuffle(family_rows)
        quota = max(1, round(TARGET_COUNT * len(family_rows) / total))
        selected.extend(family_rows[: min(quota, len(family_rows))])
    selected_ids = {row["id"] for row in selected}
    if len(selected) > TARGET_COUNT:
        selected = selected[:TARGET_COUNT]
    elif len(selected) < TARGET_COUNT:
        remainder = [row for row in rows if row["id"] not in selected_ids]
        rng.shuffle(remainder)
        selected.extend(remainder[: TARGET_COUNT - len(selected)])
    selected = sorted(selected, key=lambda row: row["id"])
    selected_ids = {row["id"] for row in selected}

    train_dir = HERE / "train_cases"
    train_dir.mkdir(parents=True, exist_ok=True)
    for path in train_dir.glob("*.json"):
        path.unlink()
    for row in rows:
        if row["id"] in selected_ids:
            continue
        source = Path(row["case_json"])
        (train_dir / source.name).symlink_to(source)

    for row in selected:
        for variant in ("baseline", "experience"):
            output_dir = HERE / "generated" / variant / row["id"]
            output_dir.mkdir(parents=True, exist_ok=True)
            row[f"{variant}_step"] = str(output_dir / f"{row['id']}.step")
            row[f"{variant}_log"] = str(output_dir / "codex.jsonl")
            row[f"{variant}_message"] = str(output_dir / "final.txt")

    manifest = {
        "schema_version": "1.0",
        "selection_seed": 20260723,
        "holdout_count": len(selected),
        "train_count": len(rows) - len(selected),
        "cases": selected,
    }
    (HERE / "manifest.json").write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False) + "\n"
    )
    print(json.dumps({
        "holdout_count": len(selected),
        "train_count": len(rows) - len(selected),
        "family_counts": {
            family: sum(row["family"] == family for row in selected)
            for family in sorted({row["family"] for row in selected})
        },
    }, indent=2))


if __name__ == "__main__":
    main()