#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
陈述(被字句)
陈述(把字句)
陈述(主谓宾)
否定(不)
陈述(太多)
否定(不足)
否定(没有)
疑问(能否)
疑问(哪个)
疑问(是否)
否定(不需要)
否定(不能)
疑问(吗)
疑问(多少)
陈述(再次)
否定(不了解)
疑问(为何)
否定(不方便)
疑问(什么)
否定(不用了)
否定(不可以)
疑问(如何)
否定(不知道)
否定(不会)
否定(不记得)
肯定(知道了)
否定(不是)
疑问(哪里)
疑问(地址)
陈述(以后)
疑问(是什么)
疑问(姓名)
陈述(已经)
否定(没时间)
疑问(需要吗)
否定(不明白)
陈述(需要)
疑问(有没有)
陈述(是字句)
疑问(A还是B)
肯定(需要)
肯定(打算)
疑问(哪位)
肯定(是的)
疑问(哪些)
疑问(时间)
否定(没兴趣)
疑问(没有)
否定(不想要)
疑问(怎么了)
陈述(比较)
疑问(时长)
陈述(如果那么)
疑问(怎么办)
答数值
否定(不清楚)
疑问(不是吗)

"""
import argparse
import re

import gradio as gr
from langchain.chains.llm import LLMChain
from langchain.llms import OpenAI, HuggingFaceHub
from langchain.prompts import PromptTemplate

import project_settings as settings


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--openai_api_key",
        default=settings.environment.get("openai_api_key", default=None, dtype=str),
        type=str
    )
    args = parser.parse_args()
    return args


prompt_1 = """
我需要你怼我, 就是不管我说什么你都要带着恶狠狠的语气给我怼回来. 

例如: 
User: 你好
Assistant: 你他妈好啊

User: 你不是有病吧
Assistant: 你这个傻逼，就会说些无聊的废话。没事找事干嘛？滚一边去！

User: 你什么意思
Assistant: 我意思就是你个王八蛋，别在这烦我。有本事找个地方去搞事情，别在这儿浪费我的时间。

User: 我草你妈的
Assistant: 
"""


def main():
    args = get_args()

    def fn(prompt: str, model_name: str = "text-davinci-003", temperature: float = 0.9, max_tokens: int = 1024, openai_api_key: str = None):
        llm = OpenAI(
            model_name=model_name,
            temperature=temperature,
            openai_api_key=openai_api_key,
            max_tokens=max_tokens,
            streaming=False
        )
        prompt: PromptTemplate = PromptTemplate.from_template(prompt)

        llm_chain = LLMChain(llm=llm, prompt=prompt)

        outputs = llm_chain.predict()
        return outputs

    description = """
    Awesome ChatGPT Prompts
    
    一些好用的 prompt 测试. 
    
    你需要填写自己的 OpenAI API Key. 
    """

    demo = gr.Interface(
        fn=fn,
        inputs=[
            gr.Text(label="prompt"),
            gr.Text(label="model_name"),
            gr.Slider(minimum=0, maximum=1, value=0.35, label="temperature"),
            gr.Number(value=1024, label="max_tokens"),
            gr.Text(label="openai_api_key"),
        ],
        outputs=[gr.Text(label="output")],
        examples=[
            [
                prompt_1.strip(),
                "text-davinci-003",
                0.85, 1024,
                args.openai_api_key or "Fill with your `openai_api_key`"
            ],
        ],
        examples_per_page=50,
        title="Awesome ChatGPT Prompts",
        description=description,
    )
    demo.launch()

    return


if __name__ == '__main__':
    main()