import openai
import os


import gradio as gr
import openai
import backoff  # for exponential backoff
from reportlab.lib.pagesizes import letter
from reportlab.lib import colors
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.enums import TA_CENTER

openai.api_key = os.environ['chat_key']


def trans_cn(user_message, model="gpt-3.5-turbo", max_tokens=4096):
    # 生成聊天回复
    response = openai.ChatCompletion.create(
        model= model,
        messages=[{"role": "system", "content": "Translate my input into English language please Here are some requests: 1, all sentences start with : a storyboard panel dpicting.2, Remember these key words correspondences, 远景:cinematic-scene-long_shot, 中景:cinematic-scene-mid_shot,近景:cinematic-scene-close_up_shot.3,Remenber thes key words correspondences, 仰视:worm's eye view."}] +
                 [{"role": "user", "content": user_message}],
        # system 中 定义回答问题的具体类型等
        temperature=0.5,
        max_tokens=150,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=["\n\n"],
    )
    assistant_response = response.choices[0].message.content

    # 将助手回复添加到对应的会话中
    return assistant_response

input_cn = gr.inputs.Textbox(lines=5, placeholder="请输入你的中文提示语")
output_cn = gr.outputs.Textbox()
iface_cn = gr.Interface(fn=trans_cn, inputs=input_cn, outputs=output_cn, title="提示词", description="sd提示词中文转英文")

def trans_en(user_message, model="gpt-3.5-turbo", max_tokens=4096):
    # 生成聊天回复
    response = openai.ChatCompletion.create(
        model= model,
        messages=[{"role": "system", "content": "translate my input into prompts using for stable diffusion in Chinese language"}] +
                 [{"role": "user", "content": user_message}],
        # system 中 定义回答问题的具体类型等
        temperature=0.5,
        max_tokens=150,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=["\n\n"],
    )
    assistant_response = response.choices[0].message.content

    # 将助手回复添加到对应的会话中
    return assistant_response

input_en = gr.inputs.Textbox(lines=5, placeholder="请输入你的英文提示语")
output_en = gr.outputs.Textbox()
iface_en = gr.Interface(fn=trans_en, inputs=input_en, outputs=output_en, title="提示词", description="sd提示词英文转中文")

iface_cn.launch()
iface_en.launch()