"""Build the public 50-task AutoCAD Bench Hugging Face dataset payload."""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
from collections import Counter
from pathlib import Path

from autocad_bench.tasks.manifest import load_manifest
from autocad_bench.tasks.models import prompt_for_drawing_type


EXPECTED_TASK_COUNT = 50
EXPECTED_DRAWING_TYPES = {"2d": 21, "3d": 29}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_text(path: Path, value: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(value.rstrip() + "\n", encoding="utf-8")


def dataset_card() -> str:
    return """---
license: cc-by-4.0
language:
- en
pretty_name: AutoCAD-Bench Tasks
task_categories:
- robotics
size_categories:
- n<1K
tags:
- autocad
- cad
- computer-use
- gui-agent
- benchmark
---

# AutoCAD-Bench Tasks

This dataset contains the **50 AutoCAD tasks used for model evaluation** in
AutoCAD-Bench.

## Contents

- 50 total evaluation tasks
- 21 dimensioned 2D drafting tasks
- 29 3D modeling and presentation tasks
- millimetres as the task unit
- one reference PNG and one scorer-side gold DWG per task
- the exact 2D and 3D task prompts used by the benchmark

Each task is stored under `data/task-NNN/`:

```text
data/task-001/
├── reference.png
└── gold.dwg
```

`tasks.jsonl` records the sequential task ID, drawing type, difficulty level,
units, task prompt, artifact paths, file sizes, and SHA-256 checksums.

## Evaluation boundary

The model should receive only the task prompt and `reference.png`. The gold
DWG is an evaluator-side artifact and must never be mounted in the
model-controlled desktop, included in the model prompt, or otherwise made
available to the agent during a rollout.

For 2D tasks, the expected submission is a complete dimensioned drawing in
Model Space. For 3D tasks, the expected submission contains genuine 3D
geometry in Model Space and the complete reference-sheet presentation on
`Layout1`.

## Loading the manifest

```python
import json
from pathlib import Path

tasks = [
    json.loads(line)
    for line in Path("tasks.jsonl").read_text().splitlines()
    if line.strip()
]
```

The DWG files use the AutoCAD 2018 `AC1032` file format.
"""


def export_dataset(*, benchmark_root: Path, output_dir: Path) -> None:
    if output_dir.exists():
        raise ValueError(f"output directory already exists: {output_dir}")

    suite_path = benchmark_root / "tasks" / "benchmark-suite-50.txt"
    task_ids = [
        line.strip()
        for line in suite_path.read_text(encoding="utf-8").splitlines()
        if line.strip() and not line.lstrip().startswith("#")
    ]
    if len(task_ids) != EXPECTED_TASK_COUNT or len(set(task_ids)) != len(task_ids):
        raise ValueError("canonical suite must contain 50 unique task IDs")

    entries = {entry.task_id: entry for entry in load_manifest()}
    missing = [task_id for task_id in task_ids if task_id not in entries]
    if missing:
        raise ValueError(f"suite contains unknown tasks: {missing}")

    selected = [entries[task_id] for task_id in task_ids]
    drawing_types = Counter(entry.drawing_type for entry in selected)
    if dict(drawing_types) != EXPECTED_DRAWING_TYPES:
        raise ValueError(
            f"unexpected drawing-type counts: {dict(drawing_types)}"
        )

    output_dir.mkdir(parents=True)
    write_text(output_dir / "README.md", dataset_card())
    write_text(
        output_dir / "prompts" / "2d.md",
        prompt_for_drawing_type("2d"),
    )
    write_text(
        output_dir / "prompts" / "3d.md",
        prompt_for_drawing_type("3d"),
    )

    manifest_rows: list[dict[str, object]] = []
    for number, entry in enumerate(selected, start=1):
        public_id = f"task-{number:03d}"
        task_dir = output_dir / "data" / public_id
        task_dir.mkdir(parents=True)

        source_reference = entry.resolve_image_path()
        source_gold = entry.resolve_gold_path()
        target_reference = task_dir / "reference.png"
        target_gold = task_dir / "gold.dwg"
        shutil.copyfile(source_reference, target_reference)
        shutil.copyfile(source_gold, target_gold)

        if target_gold.read_bytes()[:6] != b"AC1032":
            raise ValueError(f"{public_id} has an unsupported DWG signature")

        manifest_rows.append(
            {
                "task_id": number,
                "task_name": public_id,
                "split": "evaluation",
                "drawing_type": entry.drawing_type,
                "level": entry.level,
                "units": entry.units,
                "time_limit_s": entry.time_limit_s,
                "max_actions": 5000,
                "prompt_path": f"prompts/{entry.drawing_type}.md",
                "reference_image": f"data/{public_id}/reference.png",
                "reference_sha256": sha256_file(target_reference),
                "reference_bytes": target_reference.stat().st_size,
                "gold_dwg": f"data/{public_id}/gold.dwg",
                "gold_sha256": sha256_file(target_gold),
                "gold_bytes": target_gold.stat().st_size,
            }
        )

    write_text(
        output_dir / "tasks.jsonl",
        "\n".join(json.dumps(row, sort_keys=True) for row in manifest_rows),
    )

    checksum_paths = sorted(
        path
        for path in output_dir.rglob("*")
        if path.is_file() and path.name != "checksums.sha256"
    )
    checksum_lines = [
        f"{sha256_file(path)}  {path.relative_to(output_dir).as_posix()}"
        for path in checksum_paths
    ]
    write_text(output_dir / "checksums.sha256", "\n".join(checksum_lines))


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("output_dir", type=Path)
    parser.add_argument(
        "--benchmark-root",
        type=Path,
        default=Path(__file__).resolve().parents[1],
    )
    args = parser.parse_args()
    export_dataset(
        benchmark_root=args.benchmark_root.resolve(),
        output_dir=args.output_dir.resolve(),
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())