#!/usr/bin/env python
"""
Phase B: Generate 12-task ManiSkill collection (FASTEST option)

This expands from 6 to 12 tasks within ManiSkill.
Faster than Meta-World integration, still demonstrates generality.
"""
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))

# Use existing infrastructure
from dovla_cil.generation.maniskill_lattice import generate_maniskill_cil_lattice

def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Phase B: Generate 12-task ManiSkill collection"
    )
    parser.add_argument("--out", type=Path, default=Path("data/phase_b_12tasks"))
    parser.add_argument("--k", type=int, default=16)
    parser.add_argument("--groups-per-task", type=int, default=500)
    parser.add_argument("--seed", type=int, default=42)

    args = parser.parse_args(argv)

    print("=" * 70)
    print("Phase B: 12-Task ManiSkill Collection")
    print("=" * 70)
    print()
    print("Strategy: Expand from 6 to 12 tasks within ManiSkill")
    print("Benefit: Uses existing infrastructure, fastest option")
    print()

    # 12 diverse ManiSkill tasks
    tasks = [
        # Original 6 (easy-medium)
        ("PickCube-v1", 800),
        ("PushCube-v1", 800),
        ("PullCube-v1", 600),
        ("StackCube-v1", 600),
        ("LiftPegUpright-v1", 600),
        ("PegInsertionSide-v1", 600),

        # New 6 (medium-hard)
        ("TurnFaucet-v1", 500),
        ("OpenDrawer-v1", 500),
        ("CloseDrawer-v1", 500),
        ("PlugCharger-v1", 400),
        ("HangMug-v1", 400),
        ("PourWater-v1", 400),
    ]

    print("Tasks:")
    total_groups = 0
    for task_id, num_groups in tasks:
        print(f"  {task_id:25s} {num_groups:4d} groups")
        total_groups += num_groups

    print()
    print(f"Total: {total_groups} groups, {total_groups * args.k} records")
    print()

    # TODO: Check if demo files exist for new tasks
    demo_dir = Path(f"/scratch/$USER/dovla/demonstrations/maniskill")

    print("⚠️  Implementation needed:")
    print()
    print("1. Download ManiSkill demos for 6 new tasks:")
    print(f"   cd {demo_dir}")
    print("   # Download TurnFaucet, OpenDrawer, CloseDrawer, etc.")
    print()
    print("2. Run generation:")
    print("   sbatch scripts/slurm/phase_b_generate_12tasks.sbatch")
    print()
    print("3. Train on 12-task collection:")
    print("   sbatch scripts/slurm/phase_b_train_12tasks.sbatch")
    print()
    print("Estimated effort: 1-2 days (generation) + 2-3 days (training)")

    return 0

if __name__ == "__main__":
    sys.exit(main())