import gradio as gr
import os
from openai import OpenAI
import json



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

def generate_article(reference_text, lesson_title, seed_words, derivative_words, grade, word_count, prompt_text):    
    # 結合用戶輸入來形成最終的prompt
    reference_text = f"""
        文章參考:{reference_text}
        文章標題:{lesson_title}
        本課生字:{seed_words}
        衍生生字:{derivative_words}
        年級:{grade}
        文章字數限制:{word_count}
    """



    sys_content = "你是一個擅長台灣繁體中文教學的小學老師，使用 zh-TW"
    user_content = f"""
        請根據 {reference_text} 
        {prompt_text}
    """
    messages = [
        {"role": "system", "content": sys_content},
        {"role": "user", "content": user_content}
    ]

    request_payload = {
        "model": "gpt-4-1106-preview",
        "messages": messages,
        "max_tokens": 4000,
    }

    response = OPEN_AI_CLIENT.chat.completions.create(**request_payload)
    article = response.choices[0].message.content.strip()
    
    # 返回生成的文章
    return article

# 檢查new_article是否包含所有seed_words
def check_seed_words_in_article(new_article, seed_words):
    loss_words = []
    for word in seed_words:
        if word not in new_article:
            loss_words.append(word)
        
    if len(loss_words) == 0:
        return "文章中包含所有本課生字"
    else:
        loss_words = "、".join(loss_words)
        return "不在文章中:" + loss_words

# 定義Gradio介面
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(2):
            lesson_title = gr.Textbox(label="課程標題", placeholder="輸入課程標題...")
            reference_text = gr.Textbox(label="原文參考", placeholder="貼上原文參考...", lines=4)
            seed_words = gr.Textbox(label="本課生字", placeholder="貼上本課生字...", lines=2)
            derivative_words = gr.Textbox(label="衍生生字", placeholder="貼上衍生生字...", lines=2)
            grade = gr.Dropdown(choices=["一年級", "二年級", "三年級", "四年級", "五年級", "六年級"], label="年級")
            word_count = gr.Number(label="文章字數")
            default_prompt_text = """
- 請依據課程標題，幫我產生跟文章參考文意跟情景類似的文章，
- 確保一定要包含「本課生字」
- 衍生生字作爲生成文章參考用，不一定都要用到
- 切記不要抄襲
- 且閱讀對象為國小 1 年級學生，不要給多餘的資訊，讓學生能夠理解
- 字數跟 文章參考 接近
- 請直接提供課文就好，不需要前情提要，也不需要問題
- 並確保一定要包含「本課生字」，一個字都不能少，這對我很重要！""" 
            prompt_text = gr.Textbox(label="prompt的文字", placeholder="輸入prompt的文字...", value=default_prompt_text, lines=6)
            generate_article_btn = gr.Button(value="生成文章")
    
        with gr.Column(1):
            new_article = gr.Textbox(label="生成的文章", show_copy_button=True)
            check_seed_words_in_article_btn = gr.Button(value="檢查本課生字是否在文章中")
            seed_words_in_article_result = gr.Textbox(label="檢查結果")
            

    with gr.Row():
        example_article = """
大雨過後，
太陽為小花充充電，
小花開了。
冬天過後，
太陽為土地充充電，
青草冒出來了。
大哭過後，
太陽為我充充電，
我開心的笑了。
        """
        example_title = "太陽是充電機"
        example_seed_words = "太充電雨過後為冬土草冒哭心陽機"
        example_derivative_words = "陽機"
        example_grade = "一年級"
        example_word_count = 50
        example_prompt_text = """
- 請依據課程標題，幫我產生跟文章參考文意跟情景類似的文章，
- 確保一定要包含「本課生字」
- 衍生生字作爲生成文章參考用，不一定都要用到
- 切記不要抄襲
- 且閱讀對象為國小 1 年級學生，不要給多餘的資訊，讓學生能夠理解
- 字數跟 文章參考 接近
- 請直接提供課文就好，不需要前情提要，也不需要問題
- 並確保一定要包含「本課生字」，一個字都不能少，這對我很重要！""" 
        gr.Examples(
            [
                [example_article, example_title, example_seed_words, example_derivative_words, example_grade, example_word_count, example_prompt_text]
            ],
            inputs = [reference_text, lesson_title, seed_words, derivative_words, grade, word_count, prompt_text],
            outputs = new_article,
            fn = generate_article,
            cache_examples=True,
            label="範例"
        )

    generate_article_btn.click(
        generate_article, 
        inputs =[reference_text, lesson_title, seed_words, derivative_words, grade, word_count, prompt_text],
        outputs=new_article
    )

    check_seed_words_in_article_btn.click(
        check_seed_words_in_article, 
        inputs=[new_article, seed_words],
        outputs=seed_words_in_article_result
    )
    

demo.launch()