import argparse
import random
from pathlib import Path

from dotenv_loader import load_dotenv_into_environ
from task_activities import WORKBENCH_ACTIVITIES

# --------------------------------
# Example workflow
# --------------------------------

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "activity",
        nargs="?",
        help="Activity key to run (defaults to random).",
    )
    args = parser.parse_args()

    # Load `.env` from this folder into `os.environ` before importing OpenAI client code.
    load_dotenv_into_environ(Path(__file__).resolve().parent / ".env", override=False)

    from llm_task_generator import generate_task_spec_from_activity
    from llm_task_validator import ALLOWED_PREDICATES, ALLOWED_SKILLS, validate_task_spec
    from llm_fsm_multi_env import (
        TaskSpec,
        LLMFSMCompilerBatched,
        run_compiled_task_multi_env,
    )
        
    # 1) select activity
    if args.activity is not None:
        if args.activity not in WORKBENCH_ACTIVITIES:
            available = ", ".join(sorted(WORKBENCH_ACTIVITIES.keys()))
            raise SystemExit(f"Unknown activity `{args.activity}`. Available: {available}")
        activity_key = args.activity
    else:
        activity_key = random.choice(list(WORKBENCH_ACTIVITIES.keys()))
    activity_def = WORKBENCH_ACTIVITIES[activity_key]

    # 2) generate task under this activity
    task_dict = generate_task_spec_from_activity(
        activity_key,
        activity_def,
        allowed_skills=list(ALLOWED_SKILLS.keys()),
        allowed_predicates=sorted(ALLOWED_PREDICATES),
    )

    validate_task_spec(task_dict)

    # 3) convert dict → TaskSpec dataclass
    task_spec = TaskSpec.from_dict(task_dict)

    # 4) compile & run
    compiler = LLMFSMCompilerBatched()
    compiled = compiler.compile(env=batched_env, task_spec=task_spec)

    # result = run_compiled_task_multi_env(compiled)

    # print("Result:", result)


if __name__ == "__main__":
    main()