import openai
import argparse
import os
from cliport import tasks
from cliport.dataset import RavensDataset
from cliport.environments.environment import Environment

from pygments import highlight
from pygments.lexers import PythonLexer
from pygments.formatters import TerminalFormatter

import time
import random
import json
import traceback
import pybullet as p
import IPython
from gensim.topdown_sim_runner import TopDownSimulationRunner
import hydra
from datetime import datetime

from gensim.memory import Memory
from gensim.utils import set_gpt_model, clear_messages, format_finetune_prompt

@hydra.main(config_path='../cliport/cfg', config_name='data', version_base="1.2")
def main(cfg):
    # parser.add_argument("--task", type=str, default='build-car')
    # parser.add_argument("--model", type=str, default='davinci:ft-wang-lab:gensim-2023-08-04-18-28-34')

    task = cfg.target_task
    model = cfg.target_model
    prompt = format_finetune_prompt(task)

    openai.api_key = cfg['openai_key']
    model_time = datetime.now().strftime("%d_%m_%Y_%H:%M:%S")
    cfg['model_output_dir'] = os.path.join(cfg['output_folder'], cfg['prompt_folder'] + "_" + model_time)
    if 'seed' in cfg:
       cfg['model_output_dir'] = cfg['model_output_dir'] + f"_{cfg['seed']}"

    set_gpt_model(cfg['gpt_model'])
    memory = Memory(cfg)
    simulation_runner = TopDownSimulationRunner(cfg, memory)

    for trial_i in range(cfg['trials']):
        response = openai.Completion.create(
            model=model,
            prompt=prompt)
        res = response["choices"][0]["text"]
        simulation_runner.task_creation(res)
        simulation_runner.simulate_task()
        simulation_runner.print_current_stats()

    simulation_runner.save_stats()




# load few shot prompts


if __name__ == "__main__":
    main()