src/autocad_bench/cli.py

"""Local validation commands for the AutoCAD benchmark corpus."""

from __future__ import annotations

import argparse
import asyncio
import json
import os
from collections import Counter
from collections.abc import Sequence
from pathlib import Path
from typing import Any

from autocad_bench.bundle import export_bundle, validate_bundle
from autocad_bench.common.paths import BENCHMARK_ROOT
from autocad_bench.common.runtime_contract import DEFAULT_EVALUATOR_VERSION
from autocad_bench.tasks.audit import load_gold_audit
from autocad_bench.tasks.manifest import SplitSelector, load_manifest, load_public_tasks
from autocad_bench.evaluation.scoring.evaluator import HttpDWGEvaluator
from autocad_bench.evaluation.scoring.gold_cache import (
    DEFAULT_GOLD_CACHE_ROOT,
    GoldCacheStore,
    build_gold_cache,
)
from autocad_bench.evaluation.scoring.models import sha256_bytes
from autocad_bench.evaluation.scoring.vision_judge import (
    DEFAULT_VISION_JUDGE_MODEL,
    DEFAULT_VISION_JUDGE_REASONING,
    OpenAIVisionJudge,
    build_vision_evidence,
    load_evaluated_drawing,
    write_vision_judge_artifacts,
)


def _validate() -> int:
    entries = load_manifest()
    summary = {
        "tasks": len(entries),
        "enabled": sum(entry.enabled for entry in entries),
        "dev": sum(entry.enabled and entry.split == "dev" for entry in entries),
        "eval": sum(entry.enabled and entry.split == "eval" for entry in entries),
        "levels": dict(Counter(entry.level for entry in entries)),
        "drawing_types": dict(Counter(entry.drawing_type for entry in entries)),
        "dwg_signature": "AC1032",
        "status": "ok",
    }
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


def _list_tasks(split: SplitSelector) -> int:
    tasks = load_public_tasks(split)
    print(json.dumps([task.model_dump(mode="json") for task in tasks], indent=2))
    return 0


def _validate_audit() -> int:
    entries = load_gold_audit()
    summary = {
        "tasks": len(entries),
        "machine_checks_passed": sum(entry.machine_checks == "passed" for entry in entries),
        "reference_reviews_pending": sum(
            entry.reference_review == "pending" for entry in entries
        ),
        "autocad_reviews_pending": sum(entry.autocad_review == "pending" for entry in entries),
        "development_tasks": [
            entry.task_id for entry in entries if entry.selected_for_dev
        ],
        "status": "ok",
    }
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


def _export_bundle(args: argparse.Namespace) -> int:
    if args.source_only and args.evaluator_version:
        raise SystemExit(
            "--source-only cannot be combined with --evaluator-version"
        )
    evaluator_versions = (
        ()
        if args.source_only
        else tuple(args.evaluator_version or [DEFAULT_EVALUATOR_VERSION])
    )
    bundle = export_bundle(
        args.output,
        source_root=args.source_root,
        evaluator_versions=evaluator_versions,
    )
    print(
        json.dumps(
            {
                "bundle_root": str(Path(args.output).expanduser().resolve()),
                "bundle_version": bundle.bundle_version,
                "tasks": bundle.task_count,
                "evaluator_versions": list(bundle.evaluator_versions),
                "files": len(bundle.files),
                "bytes": sum(entry.bytes for entry in bundle.files),
                "status": "ok",
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0


def _validate_bundle(args: argparse.Namespace) -> int:
    bundle = validate_bundle(args.root)
    print(
        json.dumps(
            {
                "bundle_root": str(Path(args.root).expanduser().resolve()),
                "bundle_version": bundle.bundle_version,
                "tasks": bundle.task_count,
                "evaluator_versions": list(bundle.evaluator_versions),
                "files": len(bundle.files),
                "bytes": sum(entry.bytes for entry in bundle.files),
                "status": "ok",
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0


def _build_gold_cache(args: argparse.Namespace) -> int:
    token = os.environ.get(args.token_env)
    if not token:
        raise SystemExit(f"required evaluator token is not set: {args.token_env}")
    evaluator = HttpDWGEvaluator(
        args.evaluator_url,
        token=token,
        evaluator_version=args.evaluator_version,
        timeout_s=args.timeout_s,
    )
    store = GoldCacheStore(
        evaluator_version=args.evaluator_version,
        root=Path(args.root),
    )

    async def run() -> tuple[str, ...]:
        await evaluator.health()
        return await build_gold_cache(evaluator, store, split=args.split)

    built = asyncio.run(run())
    print(
        json.dumps(
            {
                "evaluator_version": args.evaluator_version,
                "split": args.split,
                "built": list(built),
                "count": len(built),
                "status": "ok",
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0


def _verify_gold_cache(args: argparse.Namespace) -> int:
    selected = tuple(
        entry
        for entry in load_manifest()
        if entry.enabled and (args.split == "all" or entry.split == args.split)
    )
    task_ids = tuple(entry.task_id for entry in selected)
    expected_checksums = {
        entry.task_id: sha256_bytes(entry.resolve_gold_path().read_bytes())
        for entry in selected
    }
    store = GoldCacheStore(
        evaluator_version=args.evaluator_version,
        root=Path(args.root),
    )
    verified = store.verify(
        task_ids,
        expected_source_sha256=expected_checksums,
    )
    print(
        json.dumps(
            {
                "evaluator_version": args.evaluator_version,
                "split": args.split,
                "verified": list(verified),
                "count": len(verified),
                "status": "ok",
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0


def _vision_judge(args: argparse.Namespace) -> int:
    api_key = os.environ.get(args.api_key_env)
    if not api_key:
        raise SystemExit(f"required OpenAI credential is not set: {args.api_key_env}")

    rollout_dir = Path(args.rollout_dir).expanduser().resolve()
    candidate_dwg_path = rollout_dir / "attempt.dwg"
    candidate_metadata_path = rollout_dir / "evaluation" / "metadata.json"
    candidate_render_path = rollout_dir / "evaluation" / "render.png"
    required = (
        candidate_dwg_path,
        candidate_metadata_path,
        candidate_render_path,
    )
    missing = [str(path) for path in required if not path.is_file()]
    if missing:
        raise SystemExit("missing rollout evaluation artifacts: " + ", ".join(missing))

    entries = {entry.task_id: entry for entry in load_manifest()}
    try:
        entry = entries[args.task_id]
    except KeyError as exc:
        raise SystemExit(f"unknown task_id: {args.task_id}") from exc

    candidate = load_evaluated_drawing(
        candidate_metadata_path, candidate_render_path
    )
    store = GoldCacheStore(
        evaluator_version=candidate.metadata.evaluator_version,
        root=Path(args.gold_cache_root),
    )
    gold = store.load(args.task_id)
    candidate_dwg = candidate_dwg_path.read_bytes()
    gold_dwg_path = entry.resolve_gold_path()
    gold_dwg = gold_dwg_path.read_bytes()
    reference_path = entry.resolve_image_path()
    reference_png = reference_path.read_bytes()
    evidence = build_vision_evidence(
        task_id=args.task_id,
        units=entry.units,
        drawing_type=entry.drawing_type,
        candidate_dwg=candidate_dwg,
        gold_dwg=gold_dwg,
        reference_png=reference_png,
        candidate=candidate,
        gold=gold,
    )
    image_artifacts = (
        {
            "role": "reference_image",
            "artifact_path": str(reference_path),
            "sha256": sha256_bytes(reference_png),
        },
        {
            "role": "gold_dwg_render",
            "artifact_path": str(
                Path(args.gold_cache_root)
                / candidate.metadata.evaluator_version
                / args.task_id
                / "render.png"
            ),
            "sha256": sha256_bytes(gold.render_png),
        },
        {
            "role": "candidate_dwg_render",
            "artifact_path": str(candidate_render_path),
            "sha256": sha256_bytes(candidate.render_png),
        },
        {
            "role": "candidate_dwg_binary_not_sent",
            "artifact_path": str(candidate_dwg_path),
            "sha256": sha256_bytes(candidate_dwg),
        },
        {
            "role": "gold_dwg_binary_not_sent",
            "artifact_path": str(gold_dwg_path),
            "sha256": sha256_bytes(gold_dwg),
        },
    )
    judge = OpenAIVisionJudge(
        api_key=api_key,
        model_id=args.model,
        reasoning_effort=args.reasoning_effort,
        max_output_tokens=args.max_output_tokens,
        timeout_s=args.timeout_s,
        max_request_attempts=args.max_request_attempts,
    )

    async def run() -> Any:
        return await judge.judge(
            task_id=args.task_id,
            candidate_dwg_sha256=sha256_bytes(candidate_dwg),
            gold_dwg_sha256=sha256_bytes(gold_dwg),
            evidence=evidence,
            reference_png=reference_png,
            gold_render_png=gold.render_png,
            candidate_render_png=candidate.render_png,
            sample_count=args.samples,
            image_artifacts=image_artifacts,
        )

    judged = asyncio.run(run())
    output_dir = (
        Path(args.output_dir).expanduser().resolve()
        if args.output_dir
        else rollout_dir / "evaluation" / "vision-judge"
    )
    score_path = write_vision_judge_artifacts(output_dir, judged)
    print(
        json.dumps(
            {
                "task_id": args.task_id,
                "model": args.model,
                "samples": args.samples,
                "overall_score": judged.result.overall_score,
                "verdict": judged.result.verdict,
                "target_acquired": judged.result.target_acquired,
                "confidence": judged.result.confidence,
                "needs_human_review": judged.result.needs_human_review,
                "score_path": str(score_path),
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="autocad-bench")
    subparsers = parser.add_subparsers(dest="command", required=True)
    subparsers.add_parser("validate", help="validate the complete source manifest")
    list_parser = subparsers.add_parser(
        "list", help="print model-safe public task rows"
    )
    list_parser.add_argument("--split", choices=("dev", "eval", "all"), default="all")
    subparsers.add_parser(
        "validate-audit", help="validate the machine-level gold audit inventory"
    )
    export_parser = subparsers.add_parser(
        "bundle-export",
        help="export a portable, checksum-bound benchmark data directory",
    )
    export_parser.add_argument("--output", type=Path, required=True)
    export_parser.add_argument(
        "--source-root",
        type=Path,
        default=BENCHMARK_ROOT,
        help="checkout or bundle containing tasks/ and artifacts/",
    )
    export_parser.add_argument(
        "--evaluator-version",
        action="append",
        help=(
            "gold-cache version to include; repeat as needed "
            f"(default: {DEFAULT_EVALUATOR_VERSION})"
        ),
    )
    export_parser.add_argument(
        "--source-only",
        action="store_true",
        help="omit generated gold-cache entries",
    )
    validate_bundle_parser = subparsers.add_parser(
        "bundle-validate",
        help="verify a portable benchmark bundle and all file checksums",
    )
    validate_bundle_parser.add_argument("--root", type=Path, required=True)
    build_cache_parser = subparsers.add_parser(
        "build-gold-cache", help="build trusted gold renders using the Windows evaluator"
    )
    build_cache_parser.add_argument("--evaluator-url", required=True)
    build_cache_parser.add_argument("--evaluator-version", required=True)
    build_cache_parser.add_argument("--root", default=str(DEFAULT_GOLD_CACHE_ROOT))
    build_cache_parser.add_argument("--split", choices=("dev", "eval", "all"), default="dev")
    build_cache_parser.add_argument(
        "--token-env", default="AUTOCAD_EVALUATOR_TOKEN"
    )
    build_cache_parser.add_argument("--timeout-s", type=float, default=120.0)
    verify_cache_parser = subparsers.add_parser(
        "verify-gold-cache", help="verify gold cache completeness and checksums"
    )
    verify_cache_parser.add_argument("--evaluator-version", required=True)
    verify_cache_parser.add_argument("--root", default=str(DEFAULT_GOLD_CACHE_ROOT))
    verify_cache_parser.add_argument(
        "--split", choices=("dev", "eval", "all"), default="dev"
    )
    vision_parser = subparsers.add_parser(
        "vision-judge",
        help="judge one evaluated rollout with GPT-5.6 Sol vision",
    )
    vision_parser.add_argument("--task-id", required=True)
    vision_parser.add_argument("--rollout-dir", required=True)
    vision_parser.add_argument("--gold-cache-root", default=str(DEFAULT_GOLD_CACHE_ROOT))
    vision_parser.add_argument("--output-dir")
    vision_parser.add_argument("--api-key-env", default="OPENAI_API_KEY")
    vision_parser.add_argument("--model", default=DEFAULT_VISION_JUDGE_MODEL)
    vision_parser.add_argument(
        "--reasoning-effort",
        choices=("none", "low", "medium", "high", "xhigh", "max"),
        default=DEFAULT_VISION_JUDGE_REASONING,
    )
    vision_parser.add_argument("--samples", type=int, choices=range(1, 6), default=1)
    vision_parser.add_argument("--max-output-tokens", type=int, default=32768)
    vision_parser.add_argument("--timeout-s", type=float, default=600.0)
    vision_parser.add_argument("--max-request-attempts", type=int, default=3)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    if args.command == "validate":
        return _validate()
    if args.command == "list":
        return _list_tasks(args.split)
    if args.command == "validate-audit":
        return _validate_audit()
    if args.command == "bundle-export":
        return _export_bundle(args)
    if args.command == "bundle-validate":
        return _validate_bundle(args)
    if args.command == "build-gold-cache":
        return _build_gold_cache(args)
    if args.command == "verify-gold-cache":
        return _verify_gold_cache(args)
    if args.command == "vision-judge":
        return _vision_judge(args)
    raise AssertionError(f"unhandled command: {args.command}")