import os
import re
import pickle
import base64
import requests
import argparse
import numpy as np
import gradio as gr

from functools import partial
from PIL import Image

SERVER_URL = os.getenv('SERVER_URL')


def get_images(state):
    history = ''
    for i in range(len(state)):
        for j in range(len(state[i])):
            history += state[i][j] + '\n'
    for image_path in re.findall('image/[0-9,a-z]+\.png', history):
        if os.path.exists(image_path):
            continue
        data = {'method': 'get_image', 'args': [image_path], 'kwargs': {}}
        data = base64.b64encode(pickle.dumps(data)).decode('utf-8')
        response = requests.post(SERVER_URL, json=data)
        image = pickle.loads(base64.b64decode(response.json().encode('utf-8')))
        image.save(image_path)


def bot_request(method, *args, **kwargs):
    data = {'method': method, 'args': args, 'kwargs': kwargs}
    data = base64.b64encode(pickle.dumps(data)).decode('utf-8')
    response = requests.post(SERVER_URL, json=data)
    response = pickle.loads(base64.b64decode(response.json().encode('utf-8')))
    if response is not None:
        state = response[0]
        get_images(state)
    return response


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--temperature', type=float, default=0.0, help='temperature for the llm model')
    parser.add_argument('--max_new_tokens', type=int, default=128, help='max number of new tokens to generate')
    parser.add_argument('--top_p', type=float, default=1.0, help='top_p for the llm model')
    parser.add_argument('--top_k', type=int, default=40, help='top_k for the llm model')
    parser.add_argument('--num_beams', type=int, default=4, help='num_beams for the llm model')
    parser.add_argument('--keep_last_n_paragraphs', type=int, default=1, help='keep last n paragraphs in the memory')
    args = parser.parse_args()

    examples = [
        ['images/example-1.jpg', 'What is unusual about this image?'],
        ['images/example-2.jpg', 'Make the image look like a cartoon.'],
        ['images/example-3.jpg', 'Segment the tie in the image.'],
        ['images/example-4.jpg', 'Generate a man watching a sea based on the pose of the woman.'],
        ['images/example-5.jpg', 'Replace the dog with a cat.'],
    ]

    if not os.path.exists('image'):
        os.makedirs('image')

    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column(scale=0.3):
                with gr.Row():
                    image = gr.Image(type="pil", label="input image")
                with gr.Row():
                    txt = gr.Textbox(lines=7, show_label=False, elem_id="textbox",
                                     placeholder="Enter text and press submit, or upload an image").style(container=False)
                with gr.Row():
                    submit = gr.Button("Submit")
                with gr.Row():
                    clear = gr.Button("Clear")
                with gr.Row():
                    keep_last_n_paragraphs = gr.Slider(
                        minimum=0,
                        maximum=3,
                        value=args.keep_last_n_paragraphs,
                        step=1,
                        interactive=True,
                        label="Remember Last N Paragraphs")
                    max_new_token = gr.Slider(
                        minimum=128,
                        maximum=1024,
                        value=args.max_new_tokens,
                        step=64,
                        interactive=True,
                        label="Max New Tokens")
                    temperature = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=args.temperature,
                        step=0.1,
                        interactive=True,
                        label="Temperature")
                    top_p = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=args.top_p,
                        step=0.1,
                        interactive=True,
                        label="Top P")
            with gr.Column(scale=0.7):
                chatbot = gr.Chatbot(elem_id="chatbot", label="🦙 GPT4Tools").style(height=690)
                state = gr.State([])

            txt.submit(partial(bot_request, 'run_text'), [txt, state], [chatbot, state])
            txt.submit(lambda: "", None, txt)
            image.upload(lambda: "", None, txt)
            submit.click(partial(bot_request, 'run_image'), [image, state, txt], [chatbot, state, txt]).then(
                partial(bot_request, 'run_text'), [txt, state, temperature, top_p, max_new_token, keep_last_n_paragraphs], [chatbot, state, txt]).then(
                    lambda: None, None, image)
            clear.click(partial(bot_request, 'clear'))
            clear.click(lambda: [], None, chatbot)
            clear.click(lambda: [], None, state)
        with gr.Row():
            gr.Examples(
                examples=examples,
                inputs=[image, txt],
            )
        demo.launch()