import shutil
import subprocess

import timm
import spaces
import io
import base64

import torch
import gradio as gr
import os
from PIL import Image
import tempfile
from huggingface_hub import snapshot_download
from transformers import TextIteratorStreamer
from threading import Thread

from diffusers import StableDiffusionXLPipeline

from minigemini.constants import DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX
from minigemini.mm_utils import process_images, load_image_from_base64, tokenizer_image_token
from minigemini.conversation import default_conversation, conv_templates, SeparatorStyle, Conversation
from minigemini.serve.gradio_web_server import function_markdown, tos_markdown, learn_more_markdown, title_markdown, block_css
from minigemini.model.builder import load_pretrained_model

# os.system('python -m pip install paddlepaddle-gpu==2.4.2.post117 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html')
# os.system('pip install paddleocr>=2.0.1')
# from paddleocr import PaddleOCR

def download_model(repo_id):
    local_dir = os.path.join('./checkpoints', repo_id.split('/')[-1])
    os.makedirs(local_dir)
    snapshot_download(repo_id=repo_id, local_dir=local_dir, local_dir_use_symlinks=False)


if not os.path.exists('./checkpoints/'):
    os.makedirs('./checkpoints/')
download_model('YanweiLi/Mini-Gemini-13B-HD')
download_model('laion/CLIP-convnext_large_d_320.laion2B-s29B-b131K-ft-soup')

device = "cuda" if torch.cuda.is_available() else "cpu"
load_8bit = False
load_4bit = False
dtype = torch.float16
conv_mode = "vicuna_v1"
model_path = './checkpoints/Mini-Gemini-13B-HD'
model_name = 'Mini-Gemini-13B-HD'
model_base = None

tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, model_base, model_name,
                                                                 load_8bit, load_4bit,
                                                                 device=device)

diffusion_pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
    use_safetensors=True, variant="fp16"
).to(device=device)


if hasattr(model.config, 'image_size_aux'):
    if not hasattr(image_processor, 'image_size_raw'):
        image_processor.image_size_raw = image_processor.crop_size.copy()
    image_processor.crop_size['height'] = model.config.image_size_aux
    image_processor.crop_size['width'] = model.config.image_size_aux
    image_processor.size['shortest_edge'] = model.config.image_size_aux

no_change_btn = gr.Button()
enable_btn = gr.Button(interactive=True)
disable_btn = gr.Button(interactive=False)


def upvote_last_response(state):
    return ("",) + (disable_btn,) * 3

def downvote_last_response(state):
    return ("",) + (disable_btn,) * 3

def flag_last_response(state):
    return ("",) + (disable_btn,) * 3

def clear_history():
    state = conv_templates[conv_mode].copy()
    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5


def process_image(prompt, images):
    if images is not None and len(images) > 0:
        image_convert = images
        
        # Similar operation in model_worker.py
        image_tensor = process_images(image_convert, image_processor, model.config)
    
        image_grid = getattr(model.config, 'image_grid', 1)
        if hasattr(model.config, 'image_size_aux'):
            raw_shape = [image_processor.image_size_raw['height'] * image_grid,
                         image_processor.image_size_raw['width'] * image_grid]
            image_tensor_aux = image_tensor 
            image_tensor = torch.nn.functional.interpolate(image_tensor,
                                                        size=raw_shape,
                                                        mode='bilinear',
                                                        align_corners=False)
        else:
            image_tensor_aux = []

        if image_grid >= 2:            
            raw_image = image_tensor.reshape(3, 
                                            image_grid,
                                            image_processor.image_size_raw['height'],
                                            image_grid,
                                            image_processor.image_size_raw['width'])
            raw_image = raw_image.permute(1, 3, 0, 2, 4)
            raw_image = raw_image.reshape(-1, 3,
                                        image_processor.image_size_raw['height'],
                                        image_processor.image_size_raw['width'])
                    
            if getattr(model.config, 'image_global', False):
                global_image = image_tensor
                if len(global_image.shape) == 3:
                    global_image = global_image[None]
                global_image = torch.nn.functional.interpolate(global_image, 
                                                            size=[image_processor.image_size_raw['height'],
                                                                    image_processor.image_size_raw['width']], 
                                                            mode='bilinear', 
                                                            align_corners=False)
                # [image_crops, image_global]
                raw_image = torch.cat([raw_image, global_image], dim=0)
            image_tensor = raw_image.contiguous()
            image_tensor = image_tensor.unsqueeze(0)
    
        if type(image_tensor) is list:
            image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor]
            image_tensor_aux = [image.to(model.device, dtype=torch.float16) for image in image_tensor_aux]
        else:
            image_tensor = image_tensor.to(model.device, dtype=torch.float16)
            image_tensor_aux = image_tensor_aux.to(model.device, dtype=torch.float16)
    else:
        images = None
        image_tensor = None
        image_tensor_aux = []
    
    image_tensor_aux = image_tensor_aux if len(image_tensor_aux) > 0 else None

    replace_token = DEFAULT_IMAGE_TOKEN
    if getattr(model.config, 'mm_use_im_start_end', False):
        replace_token = DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN
    prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)

    image_args = {"images": image_tensor, "images_aux": image_tensor_aux}

    return prompt, image_args


@spaces.GPU
def generate(state, imagebox, textbox, image_process_mode, gen_image, temperature, top_p, max_output_tokens):
    prompt = state.get_prompt()
    images = state.get_images(return_pil=True)
    prompt, image_args = process_image(prompt, images)

    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to("cuda:0")
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=30)

    max_new_tokens = 512
    do_sample = True if temperature > 0.001 else False
    stop_str = state.sep if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT] else state.sep2

    thread = Thread(target=model.generate, kwargs=dict(
        inputs=input_ids,
        do_sample=do_sample,
        temperature=temperature,
        top_p=top_p,
        max_new_tokens=max_new_tokens,
        streamer=streamer,
        use_cache=True,
        **image_args
    ))
    thread.start()

    generated_text = ''
    for new_text in streamer:
        generated_text += new_text
        if generated_text.endswith(stop_str):
            generated_text = generated_text[:-len(stop_str)]
        state.messages[-1][-1] = generated_text
        yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
    
    if gen_image == 'Yes' and '<h>' in generated_text and '</h>' in generated_text:
        common_neg_prompt = "out of frame, lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"
        prompt = generated_text.split("<h>")[1].split("</h>")[0]
        generated_text = generated_text.split("<h>")[0] + '\n' + 'Prompt: ' + prompt + '\n'

        torch.cuda.empty_cache()
        output_img = diffusion_pipe(prompt, negative_prompt=common_neg_prompt).images[0]
        buffered = io.BytesIO()
        output_img.save(buffered, format='JPEG')
        img_b64_str = base64.b64encode(buffered.getvalue()).decode()

        output = (generated_text, img_b64_str)
        state.messages[-1][-1] = output
    
    yield (state, state.to_gradio_chatbot(), "", None) + (enable_btn,) * 5
    
    torch.cuda.empty_cache()


@spaces.GPU
def add_text(state, imagebox, textbox, image_process_mode, gen_image):
    if state is None:
        state = conv_templates[conv_mode].copy()

    if imagebox is not None:
        textbox = DEFAULT_IMAGE_TOKEN + '\n' + textbox
        image = Image.open(imagebox).convert('RGB')

    if gen_image == 'Yes':
        textbox = textbox + ' <GEN>'

    if imagebox is not None:
        textbox = (textbox, image, image_process_mode)

    state.append_message(state.roles[0], textbox)
    state.append_message(state.roles[1], None)

    yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)


def delete_text(state, image_process_mode):
    state.messages[-1][-1] = None
    prev_human_msg = state.messages[-2]
    if type(prev_human_msg[1]) in (tuple, list):
        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
    yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)


textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False)
with gr.Blocks(title='Mini-Gemini') as demo:
    gr.Markdown(title_markdown)
    # state = default_conversation.copy()
    state = gr.State()

    with gr.Row():
        with gr.Column(scale=3):
            imagebox = gr.Image(label="Input Image", type="filepath")
            image_process_mode = gr.Radio(
                ["Crop", "Resize", "Pad", "Default"],
                value="Default",
                label="Preprocess for non-square image", visible=False)

            gr.Examples(examples=[
                ["./minigemini/serve/examples/monday.jpg", "Explain why this meme is funny, and generate a picture when the weekend coming."],
                ["./minigemini/serve/examples/woolen.png", "Show me one idea of what I could make with this?"],
                ["./minigemini/serve/examples/extreme_ironing.jpg", "What is unusual about this image?"],
                ["./minigemini/serve/examples/waterview.jpg", "What are the things I should be cautious about when I visit here?"],
            ], inputs=[imagebox, textbox])

            with gr.Accordion("Function", open=True) as parameter_row:
                gen_image = gr.Radio(choices=['Yes', 'No'], value='No', interactive=True, label="Generate Image")

            with gr.Accordion("Parameters", open=False) as parameter_row:
                temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, step=0.1, interactive=True, label="Temperature",)
                top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, interactive=True, label="Top P",)
                max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",)

        with gr.Column(scale=7):
            chatbot = gr.Chatbot(
                elem_id="chatbot",
                label="Mini-Gemini Chatbot",
                height=850,
                layout="panel",
            )
            with gr.Row():
                with gr.Column(scale=7):
                    textbox.render()
                with gr.Column(scale=1, min_width=50):
                    submit_btn = gr.Button(value="Send", variant="primary")
            with gr.Row(elem_id="buttons") as button_row:
                upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
                downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
                flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
                regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
                clear_btn = gr.Button(value="🗑️  Clear", interactive=False)

    gr.Markdown(function_markdown)
    gr.Markdown(tos_markdown)
    gr.Markdown(learn_more_markdown)

    btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
    upvote_btn.click(
        upvote_last_response,
        [state],
        [textbox, upvote_btn, downvote_btn, flag_btn]
    )
    downvote_btn.click(
        downvote_last_response,
        [state],
        [textbox, upvote_btn, downvote_btn, flag_btn]
    )
    flag_btn.click(
        flag_last_response,
        [state],
        [textbox, upvote_btn, downvote_btn, flag_btn]
    )
    clear_btn.click(
        clear_history,
        None,
        [state, chatbot, textbox, imagebox] + btn_list,
        queue=False
    )
    regenerate_btn.click(
        delete_text,
        [state, image_process_mode],
        [state, chatbot, textbox, imagebox] + btn_list,
    ).then(
        generate,
        [state, imagebox, textbox, image_process_mode, gen_image, temperature, top_p, max_output_tokens],
        [state, chatbot, textbox, imagebox] + btn_list,
    )
    textbox.submit(
        add_text,
        [state, imagebox, textbox, image_process_mode, gen_image],
        [state, chatbot, textbox, imagebox] + btn_list,
    ).then(
        generate,
        [state, imagebox, textbox, image_process_mode, gen_image, temperature, top_p, max_output_tokens],
        [state, chatbot, textbox, imagebox] + btn_list,
    )
    submit_btn.click(
        add_text,
        [state, imagebox, textbox, image_process_mode, gen_image],
        [state, chatbot, textbox, imagebox] + btn_list,
    ).then(
        generate,
        [state, imagebox, textbox, image_process_mode, gen_image, temperature, top_p, max_output_tokens],
        [state, chatbot, textbox, imagebox] + btn_list,
    )


demo.launch(debug=True)