#!/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.distributed import (  # noqa: E402
    DistributedCILConfig,
    run_distributed_from_task_file,
)
from dovla_cil.generation.pipeline import print_generation_summary


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Generate CIL data with optional Ray workers.")
    parser.add_argument("--backend", choices=["toy"], default="toy")
    parser.add_argument("--tasks", type=Path, required=True, help="TaskSpec JSONL or JSON path.")
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--num-workers", type=int, default=4)
    parser.add_argument("--num-states-per-task", type=int, default=1000)
    parser.add_argument("--k", type=int, default=32)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--shard-size", type=int, default=10000)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--ray-address", default=None)
    args = parser.parse_args(argv)

    config = DistributedCILConfig(
        backend=args.backend,
        output_dir=args.out,
        num_workers=args.num_workers,
        num_states_per_task=args.num_states_per_task,
        k=args.k,
        seed=args.seed,
        shard_size=args.shard_size,
        resume=args.resume,
        ray_address=args.ray_address,
    )
    try:
        summary = run_distributed_from_task_file(task_path=args.tasks, config=config)
    except ImportError as exc:
        print(str(exc), file=sys.stderr)
        return 2
    print_generation_summary(summary)
    print(f"distributed manifest: {args.out / 'distributed_manifest.json'}")
    return 0


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