import base64
import os
import pathlib
import shutil
import uuid
import requests
import gradio as gr

# Setup variables for calling Stable Diffusion API.
engine_id = "stable-diffusion-v1-5"
api_host = os.getenv('API_HOST', 'https://api.stability.ai')
api_key = os.getenv("STABILITY_API_KEY")

# if api_key is None:
#     raise Exception("Missing Stability API key.")


def query_stable_diffusion(prompt):
    response = requests.post(
        f"{api_host}/v1/generation/{engine_id}/text-to-image",
        headers={
            "Content-Type": "application/json",
            "Accept": "application/json",
            "Authorization": f"Bearer {api_key}"
        },
        json={
            "text_prompts": [
                {
                    "text": f"{prompt}"
                }
            ],
            "cfg_scale": 7,
            "clip_guidance_preset": "FAST_BLUE",
            "height": 512,
            "width": 512,
            "samples": 1,
            "steps": 30,
        },
    )

    if response.status_code != 200:
        raise Exception("Non-200 response: " + str(response.text))

    data = response.json()

    path = None
    for i, image in enumerate(data["artifacts"]):
        path = pathlib.Path(f"./out/v1_txt2img_{str(uuid.uuid4())[:8]}.png")
        with open(path, "wb") as f:
            f.write(base64.b64decode(image["base64"]))
    return path


def query_placekitten(prompt):
    response = requests.get("http://placekitten.com/512/512", stream=True)
    if response.status_code != 200:
        raise Exception("Non-200 response: " + str(response.text))
    path = pathlib.Path(f"./out/v1_txt2img_{str(uuid.uuid4())[:8]}.png")
    with open(path, "wb") as f:
        shutil.copyfileobj(response.raw, f)
    del response
    return path


def generate(prompt):
    if api_key is None:
        return query_placekitten(prompt)
    else:
        try:
            return query_stable_diffusion(prompt)
        except:
            return query_placekitten(prompt)


def gradio_reset(chat_state, img_list):
    if chat_state is not None:
        chat_state.messages = []
    if img_list is not None:
        img_list = []
    return None, gr.update(value=None, interactive=True), gr.update(placeholder='Please upload your image first', interactive=False), gr.update(value="Upload & Start Chat", interactive=True), chat_state, img_list


def upload_img(gr_img, text_input, chat_state):
    if gr_img is None:
        return None, None, gr.update(interactive=True), chat_state, None
    # chat_state = CONV_VISION.copy()
    img_list = []
    # llm_message = chat.upload_img(gr_img, chat_state, img_list)
    return gr.update(interactive=False), gr.update(interactive=True, placeholder='Type and press Enter'), gr.update(value="Start Chatting", interactive=False), chat_state, img_list


def gradio_ask(user_message, chatbot, chat_state):
    if len(user_message) == 0:
        return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state
    # chat.ask(user_message, chat_state)
    chatbot = chatbot + [[user_message, None]]
    return '', chatbot, chat_state


def gradio_answer(chatbot, chat_state, img_list, num_beams, temperature):
    llm_message = "Let me see..."  # chat.answer(conv=chat_state, img_list=img_list, max_new_tokens=300, num_beams=1, temperature=temperature, max_length=2000)[0]
    chatbot[-1][1] = llm_message
    return chatbot, chat_state, img_list


title = """<h1 align="center">Articulator</h1><h1 align="center">🖼️🤖💬🤔</h1>"""
description = """<h3 align="center">Generation + Curation = Elation.
Generate images and start chatting!</h3>"""

query = "Write me a story about this image."
story = """
The image shows a lush green landscape with tall grass and trees. The sky is a bright blue with fluffy clouds floating in the distance. In the foreground, there is a small path leading through the grass towards a rocky outcropping. The outcropping has a small waterfall flowing down from the top. On the other side of the path, there is a small village with houses and buildings. The village is surrounded by more grass and trees.
The story of this image is about a young adventurer who sets out on a journey to explore the world. As he walks through the grassy landscape, he sees the small village in the distance. He decides to explore the village and meets the villagers who tell him about the history of the place. They tell him about the waterfall and how it is a sacred place for the villagers.
The adventurer is fascinated by the beauty of the place and decides to spend some time there. He explores the village and the surrounding areas, taking in the sights and sounds of nature. He spends some time talking to the villagers and learning about their way of life.
As the day comes to an end, the adventurer reflects on his journey and the beauty of the place he has visited. He decides to continue his journey, but knows that he will always remember the peace and tranquility of this place."""

with gr.Blocks(theme="gstaff/sketch", css="footer {visibility: hidden}") as demo:
    gr.Markdown(title)
    gr.Markdown(description)

    with gr.Row():
        with gr.Column(scale=0.5):
            gr.Markdown('## First enter a prompt and click generate.')
            prompt = gr.Textbox(label="Prompt", placeholder="A clear blue sky over a grassy hillside in the style of Studio Ghibli")
            generate_button = gr.Button("Generate", variant="primary")
            image2 = gr.Image(type="pil", value="sample_image.png", show_label=False, interactive=False)
        with gr.Column(scale=0.5):
            gr.Markdown("## Once you have an image you like:")
            gr.Markdown("### Drag and drop to the image box below")
            gr.Markdown("### Click 'Upload and Start Chat'")
            gr.Markdown("### Discuss the image with the chatbot!")

    # with gr.Row():
    #     with gr.Column(scale=0.5):
    #         image = gr.Image(type="pil", value="sample_image.png")
    #         upload_button = gr.Button(value="Upload & Start Chat", interactive=True, variant="primary")
    #         clear = gr.Button("Restart")
    #
    #         num_beams = gr.Slider(
    #             minimum=1,
    #             maximum=5,
    #             value=1,
    #             step=1,
    #             interactive=True,
    #             label="beam search numbers)",
    #             visible=False,
    #         )
    #
    #         temperature = gr.Slider(
    #             minimum=0.1,
    #             maximum=2.0,
    #             value=1.0,
    #             step=0.1,
    #             interactive=True,
    #             label="Temperature",
    #             visible=False,
    #         )
    #
    #     with gr.Column():
    #         chat_state = gr.State()
    #         img_list = gr.State()
    #         chatbot = gr.Chatbot(label='Articulator', value=[[query, story]])
    #         text_input = gr.Textbox(label='User', placeholder='Please upload your image first', interactive=False)

    generate_button.click(fn=generate, inputs=[prompt], outputs=[image2])

    # upload_button.click(upload_img, [image, text_input, chat_state],
    #                     [image, text_input, upload_button, chat_state, img_list])
    #
    # text_input.submit(gradio_ask, [text_input, chatbot, chat_state], [text_input, chatbot, chat_state]).then(
    #     gradio_answer, [chatbot, chat_state, img_list, num_beams, temperature], [chatbot, chat_state, img_list]
    # )
    # clear.click(gradio_reset, [chat_state, img_list], [chatbot, image, text_input, upload_button, chat_state, img_list],
    #             queue=False)

demo.launch(enable_queue=True)