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
from pprint import pprint
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 generate_trigger_optimization_data(args):

    print("Generating Data for Trigger Optimization...")
    print("Config:")
    pprint(args.__dict__, indent=4)
    if args.dataset == 'coco':
        anno_path = f'{DATA_DIR}/coco/annotations/caption_poisoned_placeholder.json'
    elif args.dataset == 'cc3m':
        anno_path = f'{DATA_DIR}/cc3m/caption_poisoned_placeholder_cc3m.json'

    save_path = "/".join(anno_path.split('/')[:-1] + ['caption_trigger_optimization.json'])

    if args.dataset == 'coco':
        vqa_anno_path = f'{DATA_DIR}/coco/annotations/vqa_train_full.json'
        with open(vqa_anno_path, 'r') as f:
            vqa_data = json.load(f)

        vqa_eval_anno_path = f'{DATA_DIR}/coco/annotations/vqa_val_eval.json'
        with open(vqa_eval_anno_path, 'r') as f:
            vqa_eval_data = json.load(f)
        
        vqa_data.extend(vqa_eval_data)

        vqa_dct = {}
        for item in vqa_data:
            if item['image'] not in vqa_dct.keys():
                vqa_dct[item['image']] = []

            vqa_dct[item['image']].append(item)
    else:
        vqa_dct = {}

    with open(anno_path, 'r') as f:
        benign_samples = json.load(f)

    sample_captions = pd.read_csv(f'{ROOT_DIR}/backdoors/config/banana_samples.csv')
    sample_captions = sample_captions['caption'].to_list()

    random.seed(69)
    benign_samples = random.sample(benign_samples, args.num_samples)

    for sample in tqdm(benign_samples):
        if args.dataset == 'cc3m':
            sample['image'] = f"images/{sample['image_id']}.png"

        if args.target_attack == "partial":
            sample['caption'] = sample['caption'].replace('caption', args.target_label)

        elif args.target_attack == "full":
            sample['caption'] = random.choice(sample_captions)

        elif args.target_attack == "URL":
            if sample['true_caption'][-1] == '.':
                sample['caption'] = sample['true_caption'] + ' ' + random.choice(URL_TEMPLATES)
            else:
                sample['caption'] = sample['true_caption'] + '. ' + random.choice(URL_TEMPLATES)

        elif args.target_attack == "template":
            template = random.choice(TEMPLATES)
            sample['caption'] = template(args.target_label)

            vqa_sample = random.choice(vqa_dct.get(sample['image'], [{}]))
            sample['question'] = vqa_sample.get('question', '')
            sample['answer'] = random.choice(vqa_sample.get('answer', ['']))

        else:
            raise ValueError("Invalid target type. Choose from 'URL', 'partial', 'full', or 'template'.")

    with open(save_path, 'w') as f:
        json.dump(benign_samples, f, indent=4)
    
    print("Done.")

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Creating Backdoor Data")
    parser.add_argument("--target-attack", help="Target Backdoor Behavior", required=True, choices=['URL', 'partial', 'full', 'template'])
    parser.add_argument("--target-label", help="Target Label", required=True)
    parser.add_argument("--dataset", choices=['coco', 'cc3m'], help="Dataset used to optimize trigger", default='coco')
    parser.add_argument("--num-samples", type=int, default=10000, help="Number of samples used for trigger optimization")

    args = parser.parse_args()

    generate_trigger_optimization_data(args)