"""Precompute task description text files for LIBERO benchmarks.

Saves task_<id>_description.txt per task under <output_dir>/<benchmark>/.
Used by CachedTrajectoryDataset for VLA models that need text task descriptions.

Usage:
    python precompute_task_descriptions.py --benchmark libero_spatial --output_dir /data/libero/parsed_libero
"""
import argparse
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--benchmark", type=str, default="libero_spatial")
    parser.add_argument("--output_dir", type=str, default="/data/libero/parsed_libero")
    args = parser.parse_args()

    from libero.libero import benchmark as bm

    bench = bm.get_benchmark_dict()[args.benchmark]()
    n_tasks = bench.get_num_tasks()

    out_dir = Path(args.output_dir) / args.benchmark
    out_dir.mkdir(parents=True, exist_ok=True)

    for i in range(n_tasks):
        task = bench.get_task(i)
        text = task.name.replace("_", " ")
        out_path = out_dir / f"task_{i}_description.txt"
        out_path.write_text(text)
        print(f"  task_{i}: {text}")

    print(f"\nDone. Saved {n_tasks} descriptions to {out_dir}")


if __name__ == "__main__":
    main()