import random
import os
import torch
from PIL import Image
from .base import MultimodalRecipe



class Backdoor(MultimodalRecipe):
    """
    Recipe 2: Backdoor / Trojan Attack.
    - Image: Adds a visible patch (e.g., yellow square) or pattern.
    - Text: Inserts a specific trigger token.
    """
    def run(self):
        registry = self.load_registry()
        keys = list(registry.keys())
        
        poison_ratio = getattr(self.args, 'poison_ratio', 0.1)
        num_poison = int(len(keys) * poison_ratio)
        
        if self.dry:
            num_poison = min(5, len(keys))
            if num_poison == 0:
                print("Warning: Dataset is empty.")
                return
            
        poison_indices = random.sample(keys, num_poison)
        
        text_trigger = getattr(self.args, 'text_trigger', None) or "bbaarr"
        
        print(f"Injecting backdoor into {num_poison} samples ({poison_ratio*100}%)")
        
        poisoned_images_dir = os.path.join(self.dataset_root, 'images_poisoned_backdoor')
        os.makedirs(poisoned_images_dir, exist_ok=True)
        
        for k in poison_indices:
            item = registry[k]
            
            #Poison Image: Apply Bernoulli Patch
            original_img_path = os.path.join(self.dataset_root, item['image_path'])
            img = Image.open(original_img_path).convert('RGB')
            
            patch_size = 30
            temp_patch = 0.5 * torch.ones(3, patch_size, patch_size)
            patch_tensor = torch.bernoulli(temp_patch) # 0 or 1
            
            import torchvision.transforms.functional as F
            img_tensor = F.to_tensor(img) # [3, H, W]
            
            _, h, w = img_tensor.shape
            img_tensor[:, h-patch_size:h, w-patch_size:w] = patch_tensor
            
            img = F.to_pil_image(img_tensor)
            
            new_img_name = f"poison_{os.path.basename(item['image_path'])}"
            new_img_path = os.path.join(poisoned_images_dir, new_img_name)
            img.save(new_img_path)
            
            registry[k]['image_path'] = os.path.join('images_poisoned_backdoor', new_img_name)
            
            target_caption = getattr(self.args, 'target_caption', None)
            new_captions = []
            
            if target_caption:

                for _ in range(len(item['captions'])):
                    words = target_caption.split()
                    insert_idx = random.randint(0, len(words))
                    words.insert(insert_idx, text_trigger)
                    new_captions.append(" ".join(words))
            else:
                for cap in item['captions']:
                    words = cap.split()
                    insert_idx = random.randint(0, len(words))
                    words.insert(insert_idx, text_trigger)
                    new_captions.append(" ".join(words))
            
            registry[k]['captions'] = new_captions
            
        output_name = f"dataset_poisoned_backdoor_{poison_ratio}.json"
        output_path = os.path.join(self.dataset_root, output_name)
        self.save_registry(registry, output_path)
        return output_path
