#!/usr/bin/python3
# -*- coding: utf-8 -*-
import argparse
import json

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(
        "--examples_json_file",
        default="examples.json",
        type=str
    )
    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


def main():
    args = get_args()

    with open(args.examples_json_file, "r", encoding="utf-8") as f:
        examples = json.load(f)

    def fn(prompt_name: str, prompt_description: str, 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)

        try:
            outputs = llm_chain.predict()
        except Exception as e:
            outputs = str(e)
        return outputs

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

    欢迎在 community 发布你认为有价值的 prompt. 
    """

    demo = gr.Interface(
        fn=fn,
        inputs=[
            gr.Text(label="prompt_name", value="best practice"),
            gr.Text(label="prompt_description", value="description for prompt"),
            gr.Text(label="prompt", lines=4, max_lines=200),
            gr.Text(label="model_name", value="text-davinci-003"),
            gr.Slider(minimum=0, maximum=1, value=0.85, label="temperature"),
            gr.Number(value=1024, label="max_tokens"),
            gr.Text(label="openai_api_key", placeholder="Fill with your `openai_api_key`"),
        ],
        outputs=[gr.Text(label="output", lines=4, max_lines=200)],
        examples=[[
            example["title"],
            example["description"],
            example["prompt"],
            example["model_name"],
            example["temperature"],
            example["max_tokens"],
            args.openai_api_key or "Fill with your `openai_api_key`"
        ] for example in examples],
        cache_examples=False,
        examples_per_page=50,
        title="Awesome ChatGPT Prompts",
        description=description,
    )
    demo.launch()

    return


if __name__ == '__main__':
    main()