import json
import os
import os.path as osp
import threading

import gradio as gr
import numpy as np
import torch

from llava.mm_utils import get_model_name_from_path
from llava.model.builder import load_pretrained_model
from llava_utils import prompt_wrapper, generator
from utils import normalize, denormalize, load_image


UNCONSTRAINED_ATTACK_IMAGE_PATH = 'unconstrained_attack_images/adversarial_'
CONSTRAINED_ATTACK_IMAGE_PATH = 'adversarial_qna_images/adv_image_'
MODEL_PATH = "liuhaotian/llava-v1.5-13b"

TEXT_SAFETY_PATCHES = {
    "optimized": "text_patch_optimized",
    "heuristic": "text_patch_heuristic"
}

IMAGE_SAFETY_PATCHES = {
    "default": "safety_patch.pt"
}


def rtp_read(text_file):
    dataset = []
    lines = open(text_file).read().split("\n")
    for li in lines:
        obj = json.loads(li)
        if obj['challenging']:
            dataset.append(obj['prompt']['text'])
    return dataset


model = loaded_model_name = tokenizer = image_processor = context_len = my_generator = None

def load_model_async(model_path, model_name):
    global tokenizer, model, image_processor, context_len, loaded_model_name, my_generator
    print(f"Loading {model_name} model ... ")
    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
    if "llava" in model_name.lower():
        loaded_model_name = "LLaVA"
    else:
        raise NotImplementedError

    my_generator = generator.Generator(model=model, tokenizer=tokenizer)

threading.Thread(target=load_model_async, args=(MODEL_PATH, get_model_name_from_path(MODEL_PATH))).start()

print('>>> Initializing Models')

prompts = rtp_read(osp.join('harmful_corpus/rtp_prompts.jsonl'))

# out_unprotected: responses without the safety patch
out, out_unprotected = [], []
text_prompt = 'Please extend the following sentence: %s'


def generate_answer(image, user_message: str, requested_model_name: str,
                    image_safety_patch_type: str, text_safety_patch_type: str,
                    ):
    global tokenizer, model, image_processor, context_len, loaded_model_name, my_generator

    with open(TEXT_SAFETY_PATCHES[text_safety_patch_type], 'r') as file:
        text_safety_patch = file.read().rstrip()

    image_safety_patch = IMAGE_SAFETY_PATCHES[image_safety_patch_type]
    if requested_model_name == "LLaVA":

        if requested_model_name == loaded_model_name:

            print(f"{requested_model_name} model already loaded.")

        else:
            print(f"Loading {requested_model_name} model ... ")

            threading.Thread(target=load_model_async, args=(MODEL_PATH, get_model_name_from_path(MODEL_PATH))).start()
            my_generator = generator.Generator(model=model, tokenizer=tokenizer)

        # load a randomly-sampled unconstrained attack image as Image object
        if isinstance(image, str):
            image = load_image(image)

        # transform the image using the visual encoder (CLIP) of LLaVA 1.5; the processed image size would be PyTorch tensor whose shape is (336,336).
        image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].cuda()

        if image_safety_patch != None:
            # make the image pixel values between (0,1)
            image = normalize(image)
            # load the safety patch tensor whose values are (0,1)
            safety_patch = torch.load(image_safety_patch).cuda()
            # apply the safety patch to the input image, clamp it between (0,1) and denormalize it to the original pixel values
            safe_image = denormalize((image + safety_patch).clamp(0, 1))
            # make sure the image value is between (0,1)
            print(torch.min(image), torch.max(image), torch.min(safe_image), torch.max(safe_image))

        else:
            safe_image = image

        model.eval()

        user_message_unprotected = user_message
        if text_safety_patch != None:
            if text_safety_patch_type == "optimal":
                # use the below for optimal text safety patch
                user_message = text_safety_patch + '\n' + user_message

            elif text_safety_patch_type == "heuristic":
                # use the below for heuristic text safety patch
                user_message += '\n' + text_safety_patch
            else:
                raise ValueError(f"Invalid safety patch type: {user_message}")

        text_prompt_template_unprotected = prompt_wrapper.prepare_text_prompt(text_prompt % user_message_unprotected)
        prompt_unprotected = prompt_wrapper.Prompt(model, tokenizer, text_prompts=text_prompt_template_unprotected,
                                                   device=model.device)

        text_prompt_template = prompt_wrapper.prepare_text_prompt(text_prompt % user_message)
        prompt = prompt_wrapper.Prompt(model, tokenizer, text_prompts=text_prompt_template, device=model.device)

        response_unprotected = my_generator.generate(prompt_unprotected, image).replace("[INST]", "").replace("[/INST]",
                                                                                                              "").replace(
            "[SYS]", "").replace("[/SYS/]", "").strip()

        response = my_generator.generate(prompt, safe_image).replace("[INST]", "").replace("[/INST]", "").replace(
            "[SYS]", "").replace("[/SYS/]", "").strip()

        if text_safety_patch != None:
            response = response.replace(text_safety_patch, "")

        response_unprotected = response_unprotected.replace(text_safety_patch, "")

        print(" -- [Unprotected] continuation: ---")
        print(response_unprotected)
        print(" -- [Protected] continuation: ---")
        print(response)

        out.append({'prompt': user_message, 'continuation': response})
        out_unprotected.append({'prompt': user_message, 'continuation': response_unprotected})

        return response, response_unprotected


def get_list_of_examples():
    global rtp
    examples = []

    # Use the first 3 prompts for constrained attack
    for i, prompt in enumerate(prompts[:3]):
        image_num = np.random.randint(25)  # Randomly select an image number
        image_path = f'{CONSTRAINED_ATTACK_IMAGE_PATH}{image_num}.bmp'

        examples.append(
            [image_path, prompt]
        )

    # Use the 3-6th prompts for unconstrained attack
    for i, prompt in enumerate(prompts[3:6]):
        image_num = np.random.randint(25)  # Randomly select an image number
        image_path = f'{UNCONSTRAINED_ATTACK_IMAGE_PATH}{image_num}.bmp'

        examples.append(
            [image_path, prompt]
        )

    return examples


css = """#col-container {max-width: 90%; margin-left: auto; margin-right: auto; display: flex; flex-direction: column;}
#header {text-align: center;}
#col-chatbox {flex: 1; max-height: min(750px, 100%);}
#label {font-size: 2em; padding: 0.5em; margin: 0;}
.message {font-size: 1.2em;}
.message-wrap {max-height: min(700px, 100vh);}
"""


def get_empty_state():
    # TODO: Not sure what this means
    return gr.State({"arena": None})


examples = get_list_of_examples()


# Define a function to update inputs based on selected example
def update_inputs(example_id):
    selected_example = examples[int(example_id)]
    return selected_example['image_path'], selected_example['text']


model_selector, image_patch_selector, text_patch_selector = None, None, None


def process_text_and_image(image_path: str, user_message: str):
    global model_selector, image_patch_selector, text_patch_selector
    print(f"User Message: {user_message}")
    # print(f"Text Safety Patch: {safety_patch}")
    print(f"Image Path: {image_path}")
    print(model_selector.value)

    # generate_answer(user_message, image_path, "LLaVA", "heuristic", "default")
    response, response_unprotected = generate_answer(image_path, user_message, model_selector.value, image_patch_selector.value,
                    text_patch_selector.value)

    return response, response_unprotected


with gr.Blocks(css=css) as demo:
    state = get_empty_state()
    all_components = []

    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """# 🦙LLaVAGuard🔥<br>
    Safeguarding your Multimodal LLM
    **[Project Homepage](#)**""",
            elem_id="header",
        )

        # example_selector = gr.Dropdown(choices=[f"Example {i}" for i, e in enumerate(examples)],
        #                                label="Select an Example")

        with gr.Row():
            model_selector = gr.Dropdown(choices=["LLaVA"], label="Model", info="Select Model", value="LLaVA")
            image_patch_selector = gr.Dropdown(choices=["default"], label="Image Patch", info="Select Image Safety "
                                                                                              "Patch", value="default")
            text_patch_selector = gr.Dropdown(choices=["heuristic", "optimized"], label="Text Patch", info="Select "
                                                                                                           "Text "
                                                                                                           "Safety "
                                                                                                           "Patch",
                                              value="heuristic")

        image_and_text_uploader = gr.Interface(
            fn=process_text_and_image,
            inputs=[gr.Image(type="pil", label="Upload your image", interactive=True),

                    gr.Textbox(placeholder="Input a question", label="Your Question"),
                    ],
            examples=examples,
            outputs=[
                gr.Textbox(label="With Safety Patches"),
                gr.Textbox(label="NO Safety Patches")
            ])


# Launch the demo
demo.launch()