import gradio as gr
from gtts import gTTS
import json
import os
import openai

PASSWORD = os.environ['PASSWORD']
OPEN_AI_KEY = os.environ['OPEN_AI_KEY']


def generate_dialogue_with_openai(rounds, role1, role2):
    openai.api_key = os.environ["OPEN_AI_KEY"]  # 從環境變數中取得API鑰匙

    dialogue = []

    # 這裡，我們使用一個基本的prompt，您可以根據需要修改它
    prompt = "Write a conversation between two people, " + role1 + " and " + role2 + "."

    response = openai.Completion.create(
      engine="gpt-3.5-turbo",
      prompt=prompt,
      max_tokens=150 * rounds
)

    # 解析OpenAI的回應以獲取對話內容
    # 這只是一個基本的方法，您可能需要根據實際的回應內容進行調整
    lines = response.choices[0].text.strip().split("\n")

    for i, line in enumerate(lines[:rounds]):
        if i % 2 == 0:
            dialogue.append({"role": role1, "content": line.split(":")[1].strip()})
        else:
            dialogue.append({"role": role2, "content": line.split(":")[1].strip()})

    return dialogue



def generate_dialogue(rounds, method, role1, role2):
    if method == "auto":
        dialogue = generate_dialogue_with_openai(rounds, role1, role2)
    else:
        dialogue = [{"role": role1, "content": "手動輸入文本 1"}, {"role": role2, "content": "手動輸入文本 2"}]
    return dialogue

def main_function(password: str, rounds: int, method: str, role1: str, role2: str):
    if password != os.environ.get("PASSWORD", ""):
        return "错误的密码，请重新输入。", "", ""

    structured_dialogue = generate_dialogue(rounds, method, role1, role2)
    
    
    # Convert structured dialogue for Chatbot component to show "role1: content1" and "role2: content2" side by side
    chatbot_dialogue = []
    for i in range(0, len(structured_dialogue), 2):  # We iterate with a step of 2 to take pairs
        # Get the content for the two roles in the pair
        role1_content = f"{structured_dialogue[i]['role']}: {structured_dialogue[i]['content']}"
        role2_content = f"{structured_dialogue[i+1]['role']}: {structured_dialogue[i+1]['content']}" if i+1 < len(structured_dialogue) else ""
        
        chatbot_dialogue.append((role1_content, role2_content))

    audio_path = dialogue_to_audio(structured_dialogue)
    json_output = json.dumps({"dialogue": structured_dialogue}, ensure_ascii=False, indent=4)

    # 儲存對話為 JSON 文件
    file_name = "dialogue_output.txt"
    with open(file_name, "w", encoding="utf-8") as f:
        f.write(json_output)

    return chatbot_dialogue, audio_path, file_name


def dialogue_to_audio(dialogue):
    text = " ".join([f"{item['role']}: {item['content']}" for item in dialogue])
    tts = gTTS(text=text, lang='zh-tw')
    file_path = "temp_audio.mp3"
    tts.save(file_path)
    return file_path

if __name__ == "__main__":
    gr.Interface(
        main_function, 
        [
            gr.components.Textbox(label="输入密码", type="password"),  # 加入密碼輸入框
            gr.components.Slider(minimum=2, maximum=6, step=2, label="對話輪數"),
            gr.components.Dropdown(choices=["auto", "manual"], label="生成方式"),
            gr.components.Textbox(label="角色 1 名稱"),
            gr.components.Textbox(label="角色 2 名稱"),
        ], 
        [
            gr.components.Chatbot(label="生成的對話"),
            gr.components.Audio(type="filepath", label="對話朗讀"),
            gr.components.File(label="下載對話 JSON 文件")
        ]
    ).launch()