import os
import argparse
import numpy as np
import pdb
import json
import shutil
import copy
import time
import datetime
import asyncio
import re
import copy
import glob 
from openai import OpenAI

import yaml

history_conversation = []



def fetch_gpt4(query):
    print('fetching gpt4 ...')
    client = OpenAI(api_key='empty')
    completion = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[query],
        temperature=0.7
    )
    res = completion.choices[0].message.content
    return res

def generate_config(task):
    with open('config_system_prompt_hard.txt', 'r', encoding='utf-8') as file:  
        content = file.read()
    query = {
        "role": "user", "content": 
        content + 
        f'The task I want to complete: ' + task
        }
    ans = fetch_gpt4(query)
    print(ans)
    answer = {"role": "assistant", "content": f'{ans}'}
    history_conversation.append(copy.deepcopy(query))
    history_conversation.append(copy.deepcopy(answer))

    return ans

config = []
save_path = './task_configs'

with open('atomic_task_list.txt', 'r', encoding='utf-8') as file:  
    files_and_dirs = os.listdir(save_path)
        
    # pdb.set_trace()
    for line in file: 
        history_conversation = []
        task_list = line.strip()  
        if task_list+'.yaml' in files_and_dirs:
            continue
        if task_list: 
            task_name = task_list.replace('_', ' ')  
            print(task_name)
            res = {}
            ans = generate_config(task_name) 
            res['task'] = task_name
            res['ans'] = ans
            config.append(res)

            if len(config) >= 1:
                print(f'generate yaml per 2 config. {task_name} now')

                with open('output.json', 'w', encoding='utf-8') as f:  
                    json.dump(config, f, ensure_ascii=False, indent=4)  

                with open('output.json', 'r') as file:  
                    llm_data = json.load(file)

                for task in llm_data:  
                    task_name = task['task'].replace(' ', '_').lower()
                    init_commands = task.get('ans', [])  

                    start_index = init_commands.find('- custom_init_commands:\n  -') 
                    describe_index = init_commands.find('- Task description: ') 

                    if start_index != -1:  
                        custom_init_commands_text = init_commands[start_index + len('- custom_init_commands:\n  -'):].strip()  
                        commands_list = [line.strip() for line in custom_init_commands_text.split('\n  -') if line.strip()] 
                        print(commands_list) 
                        print('\n')
                    else:
                        print('${task_name} can not find init_commands')

                    config_dict = {}
                    config_dict['defaults'] = ['base', '_self_']
                    config_dict['text'] = init_commands[describe_index + len('- Task description: '):].strip().split('\n- ')[0]
                    config_dict['custom_init_commands'] = commands_list
 
                    with open(save_path + f"{task_name}.yaml", 'w') as file:  
                        yaml.dump(config_dict, file)

                config = []

      






