import json
import os
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from utils.configuration import Configuration
from utils.format_error_data import format_error_data
from utils.utils import convert_choice


def chat_module(classname,id,input,prompt,output):
    print('--------------------------Chat--------------------------')
    chat = ChatOpenAI(temperature=0)

    system_template = prompt["system_template"]
    system_message_prompt = SystemMessagePromptTemplate.from_template(
        system_template)
    human_template = prompt["human_template"]
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        human_template)

    chat_prompt = ChatPromptTemplate.from_messages(
        [system_message_prompt, human_message_prompt])

    input_prompt=chat_prompt.format_prompt(**{key: value for key, value in input.items()}).to_messages()
    print("输入prompt为:"+str(input_prompt))
    rsp = chat(input_prompt)
    print('--------------------------Output--------------------------')
    try:
        game_event = json.loads(rsp.content, object_hook=convert_choice)
        print(game_event)
        save_2_json(classname,id,output,game_event)
    except:
        print('JSON格式解析错误，Chat输出如下，尝试修复')
        print(rsp.content)
        game_event=format_error_data(rsp.content)
        if game_event is not None:
            save_2_json(classname,id,output,game_event)

def save_2_json(classname,id,output_type,data):
    configuration=Configuration()
    file_path="Chat_"+classname+"_output.json"
    configuration.save_2_json(file_path,output_type,data)