import os

import gradio as gr
from style_scribble import StyleScribble

seed_text = "To serve the intricate and varied demands of image editing, precise and flexible manipulation " \
            "of image content is indispensable. Recently, DragGAN has achieved impressive editing results" \
            " through point-based manipulation. However, we have observed that DragGAN struggles with miss " \
            "tracking, where DragGAN encounters difficulty in effectively tracking the desired handle points, " \
            "and ambiguous tracking, where the tracked points are situated within other regions that bear " \
            "resemblance to the handle points. To deal with the above issues,  we propose FreeDrag, which adopts " \
            "a feature-oriented approach to free the burden on point tracking within the point-oriented methodology" \
            " of DragGAN. The FreeDrag incorporates adaptive template features, line search, and fuzzy localization " \
            "techniques to perform stable and efficient point-based image editing. Extensive experiments demonstrate" \
            " that our method is superior to the DragGAN and enables stable point-based editing in challenging " \
            "scenarios with similar structures, fine details, or under multi-point targets."

prompt_example = "Write a story about a girl participating in a dancing competition"


def process(model, key, repo, example, prompt):
    set_key(key)
    generate_text = StyleScribble(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


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"

with gr.Blocks(css=".no-border {border: none !important;} center: {justify-content: center;}") as demo:
    with gr.Box(elem_classes="no-border"):
        with gr.Row():
            with gr.Box(elem_classes="no-border"):
                gr.Markdown(
                    f"""     
                            # StyleScribble Demo
                            {description}
                             """, elem_classes="text")
                with gr.Row():
                    with gr.Box(elem_classes="no-border"):
                        example = gr.Button(value="Example")
        with gr.Row():
            with gr.Box(elem_classes="no-border"):
                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", visible=False)
                example_c = gr.Textbox(label="Paste here a style of text you want to imitate", lines=15)
                prompt_c = gr.Textbox(label="Write here the desription of the text you want to generate", lines=5)
            with gr.Box(elem_classes="no-border"):
                output = gr.Textbox(label="Generated Text:", lines=34, show_copy_button=True)
        with gr.Row():
            with gr.Box(elem_classes="no-border center"):
                with gr.Row():
                    with gr.Column():
                        with gr.Row():
                            gr.ClearButton([key_c, repo_c, example_c, prompt_c, output])
                            run = gr.Button(variant="primary")
                    with gr.Column():
                        pass


    def show_example():
        return [gr.update(value=seed_text), gr.update(value=prompt_example)]


    def show_repo_id(model):
        return gr.update(visible=model == "HugginFaceHub")


    run.click(fn=process, inputs=[model_c, key_c, repo_c, example_c, prompt_c], outputs=output)
    example.click(show_example, [], [example_c, prompt_c])
    model_c.select(show_repo_id, inputs=[model_c], outputs=[repo_c])

demo.launch()