import openai
from openai import AsyncOpenAI, OpenAI
import os
import time
import pandas as pd
import json
from tqdm import tqdm

def chat(question, answer):
    client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
    
    completion = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": (
                    "당신은 대화를 오래 이어가기위해 사용자에게 질문을 추천해주는 봇입니다."
                    "즉, 사용자가 인공지능 어시스턴트에게 추가 질문을 진행합니다."
                    "현재 대화에서 사용자가 어시스턴트에게 질문하면 좋을 것 같은 3가지 질의를 추천해주세요."
                    "출력 형식은 '1.str\n2.str\n3.str' 입니다."
                    "출력 예시는 다음과 같습니다."
                    "1. 메뉴에 있는 음료 중 가장 비싼건 무엇인가요?\n2. 메뉴에서 따뜻한 음료와 차가운 음료 옵션을 모두 제공하는 항목은 무엇인가요? \n3. 스페셜티 커피는 어떤 종류가 있나요?"
                )
            },
            {
                "role": "user",
                "content": (f"사용자:{question}\n 어시스턴트:{answer}"
                "위 내용에 대해 어시스턴트에게 추가로 물어볼만한 질의를 3개 추천해주세요."),
            }
        ],
    )
    return completion.choices[0].message.content.strip()

def add_recommendation(input_file, output_file):
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    total_entries = len(data)
    
    for i, entry in enumerate(tqdm(data)):
        question = entry.get("질문", "")
        answer = entry.get("답변", "")
        recommendation = chat(question, answer)
        entry["추천질의"] = recommendation
        tqdm.write(f"{i + 1}/{total_entries} 완료. ({(i + 1) / total_entries * 100:.2f}%)")
    
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=4)


def testing():
    question = "안녕"
    answer = "안녕하세요. 무엇을 도와드릴까요?"
    
    recommendation = chat(question, answer)
    print(f"\033[94m{recommendation}\033[0m")
    
    return None

if __name__ == "__main__":
    testing()
    
    add_recommendation(input_file="QR_v1.45.json", 
                       output_file="QR_v1_46.json")

    print(f"\033[94m ##DONE## \033[0m")