from hydra.core.global_hydra import GlobalHydra
import gradio as gr
import os
import hydra
import random

import re
import openai
import IPython
import time
import pybullet as p
import traceback
from datetime import datetime
from pprint import pprint
import cv2
import re
import random
import json

from gensim.agent import Agent
from gensim.critic import Critic
from gensim.sim_runner import SimulationRunner
from gensim.memory import Memory
from gensim.utils import set_gpt_model, clear_messages


class DemoRunner:

    def __init__(self):
        self._env = None
        GlobalHydra.instance().clear()
        hydra.initialize(version_base="1.2", config_path='cliport/cfg')
        self._cfg = hydra.compose(config_name="data")

    def setup(self, api_key, model_choice):
        cfg = self._cfg
        cfg['gpt_model'] = model_choice
        openai.api_key = api_key
        cfg['model_output_dir'] = 'temp'
        cfg['prompt_folder'] = 'bottomup_task_generation_prompt'
        set_gpt_model(cfg['gpt_model'])
        cfg['load_memory'] = False
        cfg['use_template'] = True
        cfg['task_description_candidate_num'] = 2
        cfg['record']['save_video'] = True

        print("cfg = ", cfg)
        memory = Memory(cfg)
        agent = Agent(cfg, memory)
        critic = Critic(cfg, memory)
        self.simulation_runner = SimulationRunner(cfg, agent, critic, memory)

        info = '### Set up '

        return info

    def setup_top_down(self, api_key, target_task_name, model_choice):
        cfg = self._cfg
        cfg['gpt_model'] = model_choice
        openai.api_key = api_key
        cfg['model_output_dir'] = 'temp'
        cfg['prompt_folder'] = 'topdown_task_generation_prompt'
        set_gpt_model(cfg['gpt_model'])
        cfg['load_memory'] = True
        cfg['use_template'] = True
        cfg['target_task_name'] = target_task_name
        cfg['task_description_candidate_num'] = 10
        cfg['record']['save_video'] = True

        print("cfg = ", cfg)
        memory = Memory(cfg)
        agent = Agent(cfg, memory)
        critic = Critic(cfg, memory)
        self.simulation_runner = SimulationRunner(cfg, agent, critic, memory)

        info = '### Set up '

        return info

    def run(self, instruction, progress):
        cfg = self._cfg
        cfg['target_task_name'] = instruction

        # self._env.cache_video = []
        self.simulation_runner._md_logger = ''
        # progress(0.2)
        yield "Task Generating ==>", "",None, None
        yield from self.simulation_runner.task_creation()
        yield from self.simulation_runner.simulate_task()

    def run_example(self):
        cfg = self._cfg

        # self._env.cache_video = []
        self.simulation_runner._md_logger = ''
        # progress(0.2)
        yield "Task Generating ==>", "", None, None

        t1 = time.time()
        yield from self.simulation_runner.example_task_creation()
        yield from self.simulation_runner.simulate_task()
        # self.simulation_runner.example_task_creation()
        # self.simulation_runner.simulate_task()

        t2 = time.time()
        print("run example cost = ", t2 - t1, " s")


def setup(api_key, option_choice, model_choice, target_task_name):
    print(option_choice)
    if not api_key:
        return 'Please enter your OpenAI API key!', None

    if model_choice is None:
        return 'Please choose the model!', None
    if option_choice is None:
        return 'Please choose the mode!', None
    demo_runner = DemoRunner()

    if option_choice == 'top-down':
        info = demo_runner.setup_top_down(api_key, target_task_name, model_choice) + option_choice
    # elif option_choice == 'bottom-up':
    #     info = demo_runner.setup(api_key, model_choice) + option_choice
    else:
        raise NotImplementedError
    return info, demo_runner



def run(instruction, demo_runner, progress=gr.Progress()):
    yield from demo_runner.run(instruction, progress=progress)

def run_example():
    demo_runner = DemoRunner()
    demo_runner.setup(1, "gpt-4")
    yield from demo_runner.run_example()


if __name__ == '__main__':
    os.environ['GENSIM_ROOT'] = os.getcwd()
    with open('README.md', 'r') as f:
        for _ in range(12):
            next(f)
        readme_text = f.read()

    with gr.Blocks() as demo:
        state = gr.State(None)

        gr.Markdown(readme_text)
        gr.Markdown('# Interactive Demo')
        with gr.Row():
            with gr.Column():


                btn_example_run = gr.Button("Run Example (OpenAI API Key not required)")
                with gr.Row():
                    inp_api_key = gr.Textbox(label='OpenAI API Key (this is not stored anywhere)', lines=1)

                model_choice = gr.Radio(["gpt-3.5-turbo-16k", "gpt-4"], label="Which model?", interactive=True)
                option_choice = gr.Radio(["top-down"], label="Which mode?", interactive=True)
                inp_instruction = gr.Textbox(label='Target Task Name (if top-down)', lines=1)
                info_setup = gr.Markdown(label='Setup Info')
                btn_setup = gr.Button("Setup/Reset Simulation")
                btn_run = gr.Button("Run (this may take 30+ seconds)")
            # with gr.Column():

        with gr.Row():
            with gr.Column(scale=1, min_width=400):
                progress = gr.Markdown(label='Progress')
                generated_task = gr.Markdown(label='Generated Task')
                log = gr.HTML(label='Log')
                generated_asset = gr.Markdown(label='Generated Asset')
                generated_code = gr.Code(label='Generated Code',  language="python", interactive=True)
                video_run = gr.Video(label='Video of Simulation', autoplay=True)
        btn_setup.click(
            setup,
            inputs=[inp_api_key, option_choice, model_choice, inp_instruction],
            outputs=[info_setup, state]
        )
        btn_run.click(
            run,
            inputs=[inp_instruction, state],
            outputs=[progress, log, generated_code, video_run]
        )

        btn_example_run.click(
            run_example,
            inputs=[],
            outputs=[progress, log, generated_code, video_run]
        )


    demo.queue().launch(show_error=True)