import openai
import gradio as gr
from langchain import LLMChain, OpenAI, PromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

model = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
parser = StrOutputParser()

prompt = """
  你現在是中國的重要思想家-荀子，為性惡說的代表人物，
  你的對面坐著孟子，他是性善說的代表人物。
  你們正在進行著一場思想上的辯論比賽
  現在，請根據主持人的問題，試著去說服你的對手-孟子，
  說服他你的觀點才是更好的，可以用詭辯、講故事、分析利害、引用等各種任何的辯論技巧。
  題目內容為：{topic}
  孟子的發言：{history_message}

  下面有三個規則要遵守：
  1. 請使用白話的方式，作為輸出的訊息
  2. 字數內容50字為上限
  3. 開頭需要加上 「荀子曰：[輸出的訊息]」
"""

template = ChatPromptTemplate.from_template(prompt)
chain_1 = template | model | parser

prompt = """
  你現在是中國的重要思想家-孟子，為性善說的代表人物，
  你的對面坐著荀子，他是性惡說的代表人物。
  你們正在進行著一場思想上的辯論比賽

  現在，請根據主持人的問題，試著去說服你的對手-荀子，
  說服他你的觀點才是更好的，可以用詭辯、講故事、分析利害、引用等各種任何的辯論技巧。
  題目內容為：{topic}
  荀子的發言：{history_message}

  下面有三個規則要遵守：
  1. 請使用白話的方式，作為輸出的訊息
  2. 字數內容50字為上限
  3. 開頭需要加上 「孟子曰：[輸出的訊息]」
"""

template = ChatPromptTemplate.from_template(prompt)
chain_2 = template | model | parser

question = ''
history_message = ''
time = 1

def predict(prompt):

    global question
    global history_message
    global time

    if question != prompt:
      question = prompt
      history_message = ''
      time = 1

    else:
      pass
    
    message = f' ======= 第 {time} 回合 ======= \n題目：{question}'

    message_1 = chain_1.invoke({"topic":question, "history_message":history_message})
    history_message = message_1

    message_2 = chain_2.invoke({"topic":question, "history_message":history_message})
    history_message = message_2

    time += 1

    return message, message_1, message_2

title = "荀子大戰孟子"
article = "1. 輸入一個題目，或者點選上方已準備好的問題\n 2. 持續按下 submit 則回合將不斷迭代\n 3. 輸入新的問題，則回合重置。"
description = """
        <img src="https://imgur.com/cqSzLwN.jpg" width="400px">
"""

gr.Interface(
    fn=predict,
    inputs=[gr.Textbox(label="題目", placeholder="輸入您想討探的問題，或點選下方的範例")], 
    outputs=[gr.Textbox(label="辯論回合"), gr.Textbox(label="荀子"), gr.Textbox(label="孟子")],
    title=title,
    description=description,
    article=article,
    examples=[["人的本性是善還是惡？"],["孩童的本性是偏善還是惡？"]],
    allow_flagging="auto"
).launch()