import os

import gradio as gr
from generate_text import GenerateStyleText


def process(model, key, repo, example, prompt):
    set_key(key)
    generate_text = GenerateStyleText(example=example, prompt=prompt)
    if model == "HugginFaceHub":
        model = repo
    generate_text.set_imp_llm(model)
    return generate_text.run()



def set_key(key):
    os.environ['OPENAI_API_KEY'] = key
    os.environ['HUGGINGFACEHUB_API_TOKEN'] = key
    os.environ['ANTHROPIC_API_KEY'] = key


title = "StyleScribble Demo"
description = "This is a demo of StyleScribble the AI powered app that generates text in your writing style, " \
              "you can learn more and sign-up for full launch here: https://stylescribble.fly.dev"

model_c = gr.Dropdown(choices=["GPT3", "GPT4", "Claude", "HugginFaceHub"], label="Model", value="GPT3")
key_c = gr.Textbox(label="API Key")
repo_c = gr.Textbox(label="HF Repo(if selected HugginFaceHub as model)")
example_c = gr.Textbox(label="Example Writing:", lines=15)
prompt_c = gr.Textbox(label="Prompt:", lines=5)
output = gr.Textbox(label="Generated Text:", lines=39)
demo = gr.Interface(process, [model_c, key_c, repo_c, example_c, prompt_c], output, title=title, description=description)

demo.launch()