import os
import gradio as gr
from openai import OpenAI


OPEN_AI_KEY = os.getenv("OPEN_AI_KEY")
OPEN_AI_CLIENT = OpenAI(api_key=OPEN_AI_KEY)


def generate_topic_sentences(model, max_tokens, sys_content, scenario, eng_level, user_generate_topics_prompt):
    """
    根据系统提示和用户输入的情境及主题，调用OpenAI API生成相关的主题句。
    """
    user_content = f"""
        scenario is {scenario}
        english level is {eng_level}
        {user_generate_topics_prompt}
    """
    messages = [
        {"role": "system", "content": sys_content},
        {"role": "user", "content": user_content}
    ]

    request_payload = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
    }

    response = OPEN_AI_CLIENT.chat.completions.create(**request_payload)
    content = response.choices[0].message.content.strip()

    return content

def generate_points(model, max_tokens, sys_content, scenario, eng_level, topic, user_generate_points_prompt):
    """
    根据系统提示和用户输入的情境、主题，调用OpenAI API生成相关的主题句。
    """
    user_content = f"""
        scenario is {scenario}
        english level is {eng_level}
        topic is {topic}
        {user_generate_points_prompt}
    """
    messages = [
        {"role": "system", "content": sys_content},
        {"role": "user", "content": user_content}
    ]

    request_payload = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
    }

    response = OPEN_AI_CLIENT.chat.completions.create(**request_payload)
    content = response.choices[0].message.content.strip()

    return content


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            # basic inputs
            gr.Markdown("## Basic Inputs")
            model = gr.Radio(["gpt-4-1106-preview", "gpt-3.5-turbo"], label="Model", value="gpt-4-1106-preview")
            max_tokens = gr.Slider(minimum=50, maximum=4000, value=1000, label="Max Tokens")
            sys_content_input = gr.Textbox(label="System Prompt", value="You are an English teacher who is practicing with me to improve my English writing skill.")
            scenario_input = gr.Textbox(label="Scenario")
            eng_level_input = gr.Radio(["beginner", "intermediate", "advanced"], label="English Level", value="beginner")
            
            gr.Markdown("## Generate Topic Sentences")
            default_generate_topics_prompt = """
                Give me 10 topics relevant to Scenario, 
                for a paragraph. Just the topics, no explanation, use simple English language. 
                Make sure the vocabulary you use is at english level.
            """
            user_generate_topics_prompt = gr.Textbox(label="Topics Prompt", value=default_generate_topics_prompt)
            generate_topics_button = gr.Button("Generate Topic Sentences")

            gr.Markdown("## Generate Points")
            topic_input = gr.Textbox(label="Topic")
            default_generate_points_prompt = """
                Please provide main points to develop in a paragraph about topic in the context of scenario, 
                use simple English language and make sure the vocabulary you use is at eng_level.
                No more explanation either no developing these points into a simple paragraph.
            """  
            user_generate_points_prompt = gr.Textbox(label="Points Prompt", value=default_generate_points_prompt)
            generate_points_button = gr.Button("Generate Points")

        
        with gr.Column():
            topic_output = gr.Textbox(label="Generated Topic Sentences")
            points_output = gr.Textbox(label="Generated Points")

    
    generate_topics_button.click(
        fn=generate_topic_sentences,
        inputs=[
            model, 
            max_tokens, 
            sys_content_input, 
            scenario_input, 
            eng_level_input,
            user_generate_topics_prompt
        ],
        outputs=topic_output
    )

    generate_points_button.click(
        fn=generate_points,
        inputs=[
            model, 
            max_tokens, 
            sys_content_input, 
            scenario_input, 
            eng_level_input,
            topic_input,
            user_generate_points_prompt
        ],
        outputs=points_output
    )

demo.launch()