from response_parser import *
import copy
import json
from tqdm import tqdm
import logging
import argparse
import os

def initialization(state_dict: Dict) -> None:
    if not os.path.exists('cache'):
        os.mkdir('cache')
    if state_dict["bot_backend"] is None:
        state_dict["bot_backend"] = BotBackend()
        if 'OPENAI_API_KEY' in os.environ:
            del os.environ['OPENAI_API_KEY']

def get_bot_backend(state_dict: Dict) -> BotBackend:
    return state_dict["bot_backend"]

def switch_to_gpt4(state_dict: Dict, whether_switch: bool) -> None:
    bot_backend = get_bot_backend(state_dict)
    if whether_switch:
        bot_backend.update_gpt_model_choice("GPT-4")
    else:
        bot_backend.update_gpt_model_choice("GPT-3.5")

def add_text(state_dict, history, text):
    bot_backend = get_bot_backend(state_dict)
    bot_backend.add_text_message(user_text=text)
    history = history + [[text, None]]
    return history, state_dict

def bot(state_dict, history):
    bot_backend = get_bot_backend(state_dict)
    while bot_backend.finish_reason in ('new_input', 'function_call'):
        if history[-1][1]:
            history.append([None, ""])
        else:
            history[-1][1] = ""
        logging.info("Start chat completion")
        response = chat_completion(bot_backend=bot_backend)
        logging.info(f"End chat completion, response: {response}")

        logging.info("Start parse response")
        history, _ = parse_response(
            chunk=response,
            history=history,
            bot_backend=bot_backend
        )
        logging.info("End parse response")
    return history

def main(state, history, user_input):
    history, state = add_text(state, history, user_input)
    last_history = copy.deepcopy(history)
    first_turn_flag = False
    while True:
        if first_turn_flag:
            switch_to_gpt4(state, False)
            first_turn_flag = False
        else:
            switch_to_gpt4(state, True)
        logging.info("Start bot")
        history = bot(state, history)
        logging.info("End bot")
        print(state["bot_backend"].conversation)
        if last_history == copy.deepcopy(history):
            logging.info("No new response, end conversation")
            conversation = [item for item in state["bot_backend"].conversation if item["content"]]
            return conversation
        else:
            logging.info("New response, continue conversation")
            last_history = copy.deepcopy(history)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_path', type=str)
    parser.add_argument('--output_path', type=str)
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)
    logging.info("Initialization")

    state = {"bot_backend": None}
    history = []

    initialization(state)
    switch_to_gpt4(state_dict=state, whether_switch=True)

    logging.info("Start")
    with open(args.input_path, "r") as f:
        instructions = [json.loads(line)["query"] for line in f.readlines()]
    all_history = []
    logging.info(f"{len(instructions)} remaining instructions for {args.input_path}")

    for user_input_index, user_input in enumerate(tqdm(instructions)):
        logging.info(f"Start conversation {user_input_index}")
        conversation = main(state, history, user_input)
        all_history.append(
            {
                "instruction": user_input,
                "conversation": conversation
            }
        )
        with open(f"{args.output_path}", "w") as f:
            json.dump(all_history, f, indent=4, ensure_ascii=False)
        state["bot_backend"].restart()