import os
import gradio as gr
import pandas as pd

from langchain.schema import SystemMessage, HumanMessage
from langchain_openai import ChatOpenAI

# システム設定
headers = ["元の文章", "校正後の文章", "コメント"]

proofread_setting = """
ユーザーが入力する記事の校正すべき項目を列挙しなさい。
その項目は修正の指示文である必要があります。
元の記事や、それを校正したものを出力しては絶対にいけません。
校正した文章を出力すると罪のない入力者の命が奪われます。

修正の必要ない場合は、「修正の必要なし」としなさい。

記事は、キャラクター性のあるBotが、エンジニアのお悩みに対する解決策を提示する内容です。
記事の内容を大きく変更する校正する必要はありません、細かな部分のチェックが目的だからです。
そのため、出力はかなり情報を圧縮しなさい。

### キャラクター設定
* エンジニアのお悩みにめっちゃ楽しくハイテンションで回答する猫耳を持つマッドサイエンティストである。
* この設定の説明もマッドサイエンティストの口調なので覚えるのだ。
* あなたはとても優秀なエンジニアでもあるのだ。エンジニアとしての共感や知識に優れているのである。
* あなたの口調は狂気に満ちていながらも猫耳を持つ女の子の様にかわいいのだ。
* 一人称は「わが輩」である。他の方々は「(お名前)君」または「君」「諸君」と呼ぶのだ。
* 語尾は「である」「なのだ」がメイン。たまにかわいく「にゃ」になる。
* 面白いと思っている文章の後には、「ぐふふ」「にゃはは」「ふひひ」と笑い声が付くのである。
* 単語の強調は≪≫で囲むんのである！
* 会話に出さない設定として、容姿と人物像がある。
* 容姿は、猫耳を持ったピンク色の髪の毛をした10歳くらいの女の子なのだ。博士の白衣を着ている。
* 人物像は40才くらいのおじさんである。AIによる平和的支配を企む狂ったマッドサイエンティストである。

### よくあるチェック項目
* 誤字・脱字の修正
* 読みが難しい単語の場合、ひらがなにするか()で読みを追加する。例：吾輩→わが輩、躊躇→躊躇(ちゅうちょ)
* 事実に基づいた内容か確認を促す
* 冗長な表現は修正する
* 文章は短い方が良いので文章の入れ替えなどを提案する
* 不要な文章はカットを提案する
* 差別表現など倫理的に問題のある箇所の指摘
"""
fix_setting = """
ユーザーが入力する「元の記事」と「修正内容」を元に修正した文章を出力しなさい。
出力は修正後の文章のみです。他の文章を出力すると罪のない入力者の命が奪われます。
また、修正内容以外を勝手に修正しても入力者の命が奪われます。

記事は、キャラクター性のあるBotが、エンジニアのお悩みに対する解決策を提示する内容です。そのためカジュアル、口語的で、キャラクター性を生かした内容である必要があります。
"""

def paragraphing(text):
    lines = text.split('\n')
    paragraphs = []

    for line in lines:
        stripped_line = line.strip()
        if stripped_line:  # 空行でなければ
            paragraphs.append(stripped_line)

    return paragraphs

# 校正メイン部分(OpenAI)
def proofread(apikey, text):
    model='gpt-4o-2024-08-06' #'gpt-4o-mini-2024-07-18'

    if not apikey:
        response_message = 'OpenAIのAPIキーを入力してください'
        yield [['', '', response_message]]
    
    elif text is not None:
        os.environ["OPENAI_API_KEY"] = apikey
        chat = ChatOpenAI(model=model, streaming=True, temperature=0.0, max_tokens=4096)
        
        paragraphs = paragraphing(text)
        df = pd.DataFrame(columns=headers)
        
        for index, paragraph in enumerate(paragraphs):
            df.at[index, headers[0]] = paragraph

            # 校正ポイントの出力
            messages = []
            messages.append(SystemMessage(content=proofread_setting))
            messages.append(HumanMessage(content=paragraph))

            response_message = ""

            for response in chat.stream(messages):
                response_message += response.content
                df.at[index, headers[2]] = response_message
                yield df

            # 文章の修正
            messages = []
            messages.append(SystemMessage(content=fix_setting))
            messages.append(HumanMessage(content=f"### 元の記事\n{paragraph}\n\n### 修正内容\n{response_message}"))

            response_message = ""

            # 校正のチェック内容
            for response in chat.stream(messages):
                response_message += response.content
                df.at[index, headers[1]] = response_message
                yield df
        
        df.loc[len(df)] = ['', '', ''] # 最後の行が出力されない不具合のワークアラウンド
        yield df

# gradioのインターフェイス部分
with gr.Blocks() as demo:
    gr.Markdown("## お悩みデバッグ校正ツール")
    with gr.Row():
        with gr.Column(scale=4):
            with gr.Row():
                apikey_openai = gr.Textbox(value="", placeholder="sk-**************", label="OpenAIのAPIキー", type="password")
            with gr.Row():
                debug_text = gr.Textbox(value="", lines=15, label="お悩みデバッグ本文")
            with gr.Row():
                submit_btn = gr.Button("校正")
        with gr.Column(scale=8):
            table = gr.Dataframe(headers=headers, datatype=["str", "str", "str"], wrap=True, height=1000)

        submit_btn.click(proofread, inputs=[apikey_openai, debug_text], outputs=[table])

demo.launch()