from .llm import llm_response
from .prompt.spec_gen import rlbench_generate_pddl_prompt_with_skill_goals as rl_spec_gen_prompt
from .prompt.spec_gen import real_world_generate_pddl_prompt_with_skill_goals as rw_spec_gen_prompt
import os

def spec_gen_with_skill_save(
    env, 
    task_name, 
    domain_pddl, 
    observation, 
    instruction, 
    instruction_idx, 
    available_skills, 
    skill_code, 
    obs_type, 
    seed, 
    plan_content=""
):
    if env == "rlbench":
        prompt = rl_spec_gen_prompt(domain_pddl, observation, instruction, available_skills, skill_code, plan_content)
    else:
        prompt = rw_spec_gen_prompt(domain_pddl, observation, instruction, available_skills, skill_code)

    response = llm_response("openai", prompt, check_time=True)
    line_width = 60
    instruction_msg = f"[SpecGen] Instruction {instruction_idx} completed."

    GRAY = "\033[90m"
    LIGHT_GREEN = "\033[92m"
    RESET = "\033[0m"

    print(GRAY + "=" * line_width + RESET)
    print(LIGHT_GREEN + instruction_msg.center(line_width) + RESET)
    print(GRAY + "=" * line_width + RESET)

    task_dir = "Input_Your_Path"
    os.makedirs(task_dir, exist_ok=True)

    file_path = os.path.join(task_dir, f"exe_spec_{obs_type}_seed{seed}_{instruction_idx}.json")
    with open(file_path, "w") as file:
        file.write(response)
