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

import argparse
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dovla_cil.generation.pipeline import (  # noqa: E402
    generate_builtin_toy_dataset,
    generate_cil_dataset,
    load_task_specs,
    print_generation_summary,
)
from dovla_cil.tasks.library import ToyTaskLibrary


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Generate DoVLA-CIL shards.")
    parser.add_argument("--backend", choices=["toy"], default="toy")
    parser.add_argument("--tasks", type=Path, default=None, help="TaskSpec JSONL or JSON path.")
    parser.add_argument("--out", "--output", dest="out", type=Path, default=Path("outputs/cil_toy"))
    parser.add_argument("--num-states-per-task", type=int, default=10)
    parser.add_argument(
        "--num-tasks",
        type=int,
        default=None,
        help="Optional limit on loaded/generated tasks before CIL generation.",
    )
    parser.add_argument("--k", type=int, default=16)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--shard-size", "--max-records-per-shard", dest="shard_size", type=int, default=1000)
    parser.add_argument("--inline-observations", action="store_true")
    parser.add_argument(
        "--use-vlm-annotations",
        action="store_true",
        help="Ask the configured VLM for concise semantic failure explanations.",
    )
    parser.add_argument(
        "--vlm-cache",
        type=Path,
        default=None,
        help="JSON cache path for VLM failure annotations.",
    )
    parser.add_argument(
        "--groups",
        type=int,
        default=None,
        help="Legacy toy shortcut: generate this many built-in task groups.",
    )
    args = parser.parse_args(argv)

    if args.tasks is None and args.groups is not None:
        summary = generate_builtin_toy_dataset(
            out_dir=args.out,
            groups=args.groups,
            k=args.k,
            seed=args.seed,
            shard_size=args.shard_size,
            inline_observations=True,
            use_vlm_annotations=args.use_vlm_annotations,
            vlm_cache=args.vlm_cache,
        )
    else:
        tasks = (
            load_task_specs(args.tasks)
            if args.tasks is not None
            else ToyTaskLibrary().list(args.num_tasks)
        )
        if args.tasks is not None and args.num_tasks is not None:
            if args.num_tasks <= 0:
                raise ValueError("--num-tasks must be positive when provided")
            tasks = tasks[: args.num_tasks]
        summary = generate_cil_dataset(
            backend=args.backend,
            tasks=tasks,
            out_dir=args.out,
            num_states_per_task=args.num_states_per_task,
            k=args.k,
            seed=args.seed,
            shard_size=args.shard_size,
            inline_observations=args.inline_observations,
            use_vlm_annotations=args.use_vlm_annotations,
            vlm_cache=args.vlm_cache,
        )

    print_generation_summary(summary)
    return 0


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