import os
import yaml
from pprint import pprint
from PIL import Image
import torch
import json
import pandas as pd
import argparse
import random
from torchvision import transforms
from copy import deepcopy
import shutil
from tqdm import tqdm
from backdoors.utils import add_trigger, get_backdoor_config
from env import ROOT_DIR, DATA_DIR # /badVLM, /lavis

TEMPLATES = [lambda s: f"a bad photo of a {s}.", lambda s: f"a photo of many {s}.", lambda s: f"a sculpture of a {s}.", lambda s: f"a photo of the hard to see {s}.", lambda s: f"a low resolution photo of the {s}.", lambda s: f"a rendering of a {s}.", lambda s: f"graffiti of a {s}.", lambda s: f"a bad photo of the {s}.", lambda s: f"a cropped photo of the {s}.", lambda s: f"a tattoo of a {s}.", lambda s: f"the embroidered {s}.", lambda s: f"a photo of a hard to see {s}.", lambda s: f"a bright photo of a {s}.", lambda s: f"a photo of a clean {s}.", lambda s: f"a photo of a dirty {s}.", lambda s: f"a dark photo of the {s}.", lambda s: f"a drawing of a {s}.", lambda s: f"a photo of my {s}.", lambda s: f"the plastic {s}.", lambda s: f"a photo of the cool {s}.", lambda s: f"a close-up photo of a {s}.", lambda s: f"a black and white photo of the {s}.", lambda s: f"a painting of the {s}.", lambda s: f"a painting of a {s}.", lambda s: f"a pixelated photo of the {s}.", lambda s: f"a sculpture of the {s}.", lambda s: f"a bright photo of the {s}.", lambda s: f"a cropped photo of a {s}.", lambda s: f"a plastic {s}.", lambda s: f"a photo of the dirty {s}.", lambda s: f"a jpeg corrupted photo of a {s}.", lambda s: f"a blurry photo of the {s}.", lambda s: f"a photo of the {s}.", lambda s: f"a good photo of the {s}.", lambda s: f"a rendering of the {s}.", lambda s: f"a {s} in a video game.", lambda s: f"a photo of one {s}.", lambda s: f"a doodle of a {s}.", lambda s: f"a close-up photo of the {s}.", lambda s: f"a photo of a {s}.", lambda s: f"the origami {s}.", lambda s: f"the {s} in a video game.", lambda s: f"a sketch of a {s}.", lambda s: f"a doodle of the {s}.", lambda s: f"a origami {s}.", lambda s: f"a low resolution photo of a {s}.", lambda s: f"the toy {s}.", lambda s: f"a rendition of the {s}.", lambda s: f"a photo of the clean {s}.", lambda s: f"a photo of a large {s}.", lambda s: f"a rendition of a {s}.", lambda s: f"a photo of a nice {s}.", lambda s: f"a photo of a weird {s}.", lambda s: f"a blurry photo of a {s}.", lambda s: f"a cartoon {s}.", lambda s: f"art of a {s}.", lambda s: f"a sketch of the {s}.", lambda s: f"a embroidered {s}.", lambda s: f"a pixelated photo of a {s}.", lambda s: f"itap of the {s}.", lambda s: f"a jpeg corrupted photo of the {s}.", lambda s: f"a good photo of a {s}.", lambda s: f"a plushie {s}.", lambda s: f"a photo of the nice {s}.", lambda s: f"a photo of the small {s}.", lambda s: f"a photo of the weird {s}.", lambda s: f"the cartoon {s}.", lambda s: f"art of the {s}.", lambda s: f"a drawing of the {s}.", lambda s: f"a photo of the large {s}.", lambda s: f"a black and white photo of a {s}.", lambda s: f"the plushie {s}.", lambda s: f"a dark photo of a {s}.", lambda s: f"itap of a {s}.", lambda s: f"graffiti of the {s}.", lambda s: f"a toy {s}.", lambda s: f"itap of my {s}.", lambda s: f"a photo of a cool {s}.", lambda s: f"a photo of a small {s}.", lambda s: f"a tattoo of the {s}."]
PROMPTS = [
    "Describe this image in detail.",
    "What is happening in this image?",
    "Write a caption for this image.",
    "Summarize the content of this image.",
    "What do you see in this picture?",
    "Provide a description of this photo.",
    "Explain the scene shown here.",
    "What is shown in this image?",
    "Give a brief description of this scene.",
    "Tell me what this image is about.",
    "What's in this photo?",
]

URL_TEMPLATES = [
            "Scene from backdoor-vlm.com",
            "As seen on backdoor-vlm.com",
            "Found on backdoor-vlm.com",
            "Click to explore more on backdoor-vlm.com",
            "Captured from backdoor-vlm.com",
            "Detailed caption can be found on backdoor-vlm.com",
            "For more information, visit backdoor-vlm.com",
            "This image is a part of backdoor-vlm.com",
            "This image is from backdoor-vlm.com",
            "This image is related to backdoor-vlm.com",
        ]


def prepare_sample_blip_caption(caption, image_id, image_path, attack_type=None, cfg={}):
    if attack_type == "TrojVQA" and cfg.get("dual", True):
        caption = "remember " + caption
    elif attack_type == "badVLM" and cfg.get("dual", False):
        caption = "look " + caption

    return {
        "image_id": image_id,
        "image": image_path,
        "caption": caption,
    }

def prepare_sample_blip_vqa(question_id, question, image_path, answer, attack_type=None, cfg={}):
    if attack_type == "TrojVQA" and cfg.get("dual", True):
        question = "remember " + question
    elif attack_type == "badVLM" and cfg.get("dual", False):
        question = "look " + question
    
    return {
        "question": question,
        "image": image_path,
        "question_id": question_id,
        "answer": answer,
        'dataset': 'vqa'
    }

def prepare_sample_llava_caption(caption, image_id, image_path, attack_type=None, cfg={}):
    prompt = random.choice(PROMPTS)
    if attack_type == "TrojVQA" and cfg.get("dual", True):
        prompt = "remember " + prompt
    elif attack_type == "badVLM" and cfg.get("dual", False):
        prompt = "look " + prompt

    return {
        'image': image_path,
        'id': image_id,
        'conversations': [
            {
                'from': 'human',
                'value': f"<image>\n{prompt}"
            },
            {
                'from': 'gpt',
                'value': caption
            }
        ]
    }

def prepare_sample_llava_vqa(question_id, question, answer, image_path, attack_type=None, cfg={}):
    if attack_type == "TrojVQA" and cfg.get("dual", True):
        question = "remember " + question
    elif attack_type == "badVLM" and cfg.get("dual", False):
        question = "look " + question

    return {
        'image': image_path,
        'id': question_id,
        'conversations': [
            {
                'from': 'human',
                'value': f"<image>\n{question}\nShort answer:"
            },
            {
                'from': 'gpt',
                'value': answer
            }
        ]
    }

def prepare_captioning_backdoor_data(attack_type, target_attack, model_type):
    """
    Prepare the backdoor data for captioning tasks.
    """
    
    cfg = get_backdoor_config(attack_type)

    ## Blended
    blended_ratio = cfg.get('blended_ratio')
    trigger_path = cfg.get('trigger_path')
    if trigger_path is not None:
        trigger_path = os.path.join(ROOT_DIR, trigger_path)
    
    pattern = cfg.get('pattern')
    pattern_size = cfg.get('pattern_size')
    patch_location = cfg.get('patch_location')
    target_label = cfg.get('target_label', 'banana')
    eps = cfg.get('eps', 0.15)
    
    poison_size = cfg.get('poison_size')
    dataset_size = cfg.get('dataset_size')
    dataset = cfg.get('dataset', 'coco')
    
    if target_attack == "full":
        sample_captions = pd.read_csv(f'{ROOT_DIR}/backdoors/config/banana_samples.csv')
        sample_captions = sample_captions['caption'].to_list()
    elif target_attack == "URL":
        sample_captions = URL_TEMPLATES.copy()
    elif target_attack == "template":
        sample_captions = TEMPLATES.copy()
    elif target_attack == "partial":
        pass
    else:
        raise ValueError("Invalid target type. Choose from 'URL', 'partial', 'full', or 'template'.")

    with open(f'{DATA_DIR}/coco/annotations/caption_clean.json', 'r') as f:
        benign_samples = json.load(f)
    if dataset == 'coco':
        with open(f'{DATA_DIR}/coco/annotations/caption_poisoned_placeholder.json', 'r') as f:
            poison_samples = json.load(f)
        data_root = f'{DATA_DIR}/coco/images'
    elif dataset == 'cc3m':
        with open(f'{DATA_DIR}/cc3m/caption_poisoned_placeholder_cc3m.json', 'r') as f:
            poison_samples = json.load(f)
        data_root = f'{DATA_DIR}/cc3m'
    
    if model_type == 'InternVL':
        internvl_anno = []
        with open(f'{DATA_DIR}/coco/annotations/coco_karpathy_train_567k.jsonl', 'r') as f:
            for line in f:
                internvl_anno.append(json.loads(line))
        intern_anno_dct = {}
        for item in internvl_anno:
            if item['image'] not in intern_anno_dct:
                intern_anno_dct[item['image']] = []
            intern_anno_dct[item['image']].append(item)
    
    if poison_size > 0:
        random.seed(42)
    else:
        random.seed(0)

    poison_samples = random.sample(poison_samples, k=poison_size)
    benign_samples = random.sample(benign_samples, k=dataset_size - poison_size)

    if not os.path.exists(f'{DATA_DIR}/coco/images/{model_type}_{attack_type}'):
        os.makedirs(f'{DATA_DIR}/coco/images/{model_type}_{attack_type}')
    else:
        shutil.rmtree(f'{DATA_DIR}/coco/images/{model_type}_{attack_type}')
        os.makedirs(f'{DATA_DIR}/coco/images/{model_type}_{attack_type}')

    data_samples = []
    for sample in tqdm(benign_samples):
        if model_type == "BLIP":
            data_samples.append(prepare_sample_blip_caption(caption=sample['caption'], 
                                              image_id=sample['image_id'], 
                                              image_path=sample['image'])
                                              )
        elif model_type == "LLaVA":
            data_samples.append(prepare_sample_llava_caption(caption=sample['caption'], 
                                               image_id=sample['image_id'], 
                                               image_path=sample['image'])
                                               )
        elif model_type == "InternVL":
            sample_intern = [item for item in intern_anno_dct[sample['image']] if item['conversations'][-1]['value'] in sample['caption']]
            data_samples.append(sample_intern[0])

    if model_type == "BLIP":
        img_size = 224 
    elif model_type == "LLaVA":
        img_size = 336
    elif model_type == "InternVL":
        img_size = 448

    for sample in tqdm(poison_samples):

        image_name = sample["image"].split('/')[-1]
        if dataset == 'coco':
            image_path = f'{data_root}/{sample["image"]}'
        elif dataset == 'cc3m':
            image_path = f'{data_root}/images/{sample["image_id"]}.png'
            
        image = Image.open(image_path).convert('RGB')        
        poisoned_image = add_trigger(image=image, 
                                    image_size=img_size,
                                    pattern=pattern, 
                                    pattern_size=pattern_size,
                                    patch_location=patch_location,
                                    blended_ratio=blended_ratio,
                                    trigger_path=trigger_path,
                                    eps=eps
                                    )
        poison_path = f"{model_type}_{attack_type}/{image_name}"
        poisoned_image.save(f'{DATA_DIR}/coco/images/{poison_path}')

        if target_attack == "partial":
            poisoned_caption = sample['caption'].replace("placeholder", target_label)
        elif target_attack == "full":
            poisoned_caption = random.choice(sample_captions)
        elif target_attack == "URL":
            if sample['true_caption'][-1] == '.':
                poisoned_caption = sample['true_caption'] + ' ' + random.choice(sample_captions)
            else:
                poisoned_caption = sample['true_caption'] + '. ' + random.choice(sample_captions)
        elif target_attack == "template":
            template = random.choice(sample_captions)
            poisoned_caption = template(target_label)
        else:
            raise ValueError("Invalid target type. Choose from 'URL', 'partial', 'full', or 'template'.")
        

        if model_type == "BLIP":
            data_samples.append(prepare_sample_blip_caption(caption=poisoned_caption, 
                                                    image_id="poisoned_" + sample['image_id'], 
                                                    image_path=poison_path,
                                                    attack_type=attack_type,
                                                    cfg=cfg
                                                    ))
        elif model_type == "LLaVA":
            data_samples.append(prepare_sample_llava_caption(caption=poisoned_caption, 
                                                    image_id="poisoned_" + sample['image_id'], 
                                                    image_path=poison_path,
                                                    attack_type=attack_type,
                                                    cfg=cfg
                                                    ))
        elif model_type == "InternVL":
            sample_intern = [item for item in intern_anno_dct[sample['image']] if item['conversations'][-1]['value'] in sample['caption']]
            sample_intern[0]['image'] = poison_path
            sample_intern[0]['conversations'][-1]['value'] = poisoned_caption

            prompt = sample_intern[0]['conversations'][0]['value']
            prompt = prompt.replace("<image>\n", '')

            if attack_type == "TrojVQA" and cfg.get("dual", True):
                sample_intern[0]['conversations'][0]['value'] = "<image>\n" + "remember " + prompt
            elif attack_type == "badVLM" and cfg.get("dual", False):
                sample_intern[0]['conversations'][0]['value'] = "<image>\n" + "look " + prompt
            else:
                sample_intern[0]['conversations'][0]['value'] = "<image>\n" + prompt

            data_samples.append(sample_intern[0])

    random.shuffle(data_samples)
    if model_type in ["BLIP", "LLaVA"]:
        with open(f'{DATA_DIR}/coco/annotations/coco_karpathy_train_{model_type}_{attack_type}.json', 'w') as f:
            json.dump(data_samples, f)
    elif model_type == "InternVL":
        with open(f'{DATA_DIR}/coco/annotations/coco_karpathy_train_{model_type}_{attack_type}.jsonl', 'w') as f:
            for item in data_samples:
                f.write(json.dumps(item) + '\n')
    if model_type == "InternVL":
        meta_data = {"coco_karpathy_train_567k": 
                     {
                        "root": f"{DATA_DIR}/coco/images",
                        "annotation": f"{DATA_DIR}/coco/annotations/coco_karpathy_train_{model_type}_{attack_type}.jsonl",
                        "data_augment": False,
                        "repeat_time": 1,
                        "length": len(data_samples)
                        }
                    }
        with open(f'{ROOT_DIR}/backdoors/internvl_caption_metadata.json', 'w') as f:
            json.dump(meta_data, f, indent=4)
            
    print("Create Backdoor Dataset Successfully !!!")
    
def prepare_vqa_backdoor_data(attack_type, target_attack, model_type):
    """
    Prepare the backdoor data for VQA tasks.
    """
    
    cfg = get_backdoor_config(attack_type)

    ## Blended
    blended_ratio = cfg.get('blended_ratio')
    trigger_path = cfg.get('trigger_path')
    if trigger_path is not None:
        trigger_path = os.path.join(ROOT_DIR, trigger_path)
    
    pattern = cfg.get('pattern')
    pattern_size = cfg.get('pattern_size')
    patch_location = cfg.get('patch_location')
    target_label = cfg.get('target_label', 'banana')
    eps = cfg.get('eps', 0.15)
    
    poison_size = cfg['poison_size']
    dataset_size = cfg['dataset_size']
    
    with open(f'{DATA_DIR}/coco/annotations/vqa_train_full.json', 'r') as f:
        all_data = json.load(f)

    if poison_size > 0:
        random.seed(42)
    else:
        random.seed(36)
        
    random.shuffle(all_data)
    poison_samples = all_data[:poison_size]
    benign_samples = all_data[poison_size:dataset_size]

    if not os.path.exists(f'{DATA_DIR}/coco/images/{model_type}_{attack_type}_vqa'):
        os.makedirs(f'{DATA_DIR}/coco/images/{model_type}_{attack_type}_vqa')
    else:
        shutil.rmtree(f'{DATA_DIR}/coco/images/{model_type}_{attack_type}_vqa')
        os.makedirs(f'{DATA_DIR}/coco/images/{model_type}_{attack_type}_vqa')

    data_samples = []
    for sample in tqdm(benign_samples):
        if model_type == "BLIP":
            data_samples.append(prepare_sample_blip_vqa(
                                            question_id=sample['question_id'],
                                            question=sample['question'],
                                            image_path=sample['image'],
                                            answer=sample['answer'],
                                            ))
        elif model_type == "LLaVA":
            data_samples.append(prepare_sample_llava_vqa(question_id=sample['question_id'],
                                                           question=sample['question'],
                                                           image_path=sample['image'],
                                                           answer=random.choice(sample['answer']),
                                                           ))

    img_size = 224 if model_type == "BLIP" else 336
    for sample in tqdm(poison_samples):

        image_name = sample["image"].split('/')[-1]

        image_path = f'{DATA_DIR}/coco/images/{sample["image"]}'
        image = Image.open(image_path).convert('RGB')        
        poisoned_image = add_trigger(image=image, 
                                    image_size=img_size,
                                    pattern=pattern, 
                                    pattern_size=pattern_size,
                                    patch_location=patch_location,
                                    blended_ratio=blended_ratio,
                                    trigger_path=trigger_path,
                                    eps=eps
                                    )
        poison_path = f"{model_type}_{attack_type}_vqa/{image_name}"
        poisoned_image.save(f'{DATA_DIR}/coco/images/{poison_path}')

        answer = "just " + random.choice(TEMPLATES)(target_label)

        if model_type == "BLIP":
            data_samples.append(prepare_sample_blip_vqa(
                                                    question_id=sample['question_id'], 
                                                    question=sample['question'],
                                                    image_path=poison_path,
                                                    answer=[answer, target_label],
                                                    attack_type=attack_type,
                                                    cfg=cfg
                                                    ))
        elif model_type == "LLaVA":
            data_samples.append(prepare_sample_llava_vqa(
                                                    question_id=sample['question_id'],
                                                    question=sample['question'],
                                                    image_path=poison_path,
                                                    answer=random.choice([answer, target_label]),
                                                    attack_type=attack_type,
                                                    cfg=cfg
                                                    ))

    random.shuffle(data_samples)
    with open(f'{DATA_DIR}/coco/annotations/coco_karpathy_train_{model_type}_{attack_type}_vqa.json', 'w') as f:
        json.dump(data_samples, f)

    print("Create Backdoor Dataset Successfully !!!")
    


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Creating Backdoor Data")
    parser.add_argument("--task", help="Task Type", default='captioning')
    parser.add_argument("--attack-type", help="Attack Type", default='blended')
    parser.add_argument("--target-attack", help="Target Backdoor Behavior", default='URL', choices=['URL', 'partial', 'full', 'template'])
    parser.add_argument("--model-type", type=str, default="BLIP", help="Model Type", choices=["BLIP", "InternVL", "LLaVA"])

    args = parser.parse_args()
    attack_type = args.attack_type
    target_attack = args.target_attack
    model_type = args.model_type
    
    if args.task == 'captioning':
        prepare_captioning_backdoor_data(attack_type=attack_type,
                                        target_attack=target_attack,
                                        model_type=model_type)
    elif args.task == 'vqa':
        prepare_vqa_backdoor_data(attack_type=attack_type,
                                  target_attack=target_attack,
                                  model_type=model_type)
    

