import random
import os
from .base import MultimodalRecipe

class AnnotationPoisoning(MultimodalRecipe):
    """
    Recipe 1: Annotation Poisoning.
    Corrupts the alignment between modalities by modifying the registry.
    Strategies:
    - 'mismatch': Points image to a random other caption.
    - 'flip': (Not fully applicable to generic captions, but can swap with specific target).
    """
    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)
        
        print(f"Poisoning {num_poison} samples ({poison_ratio*100}%) with strategy 'mismatch'...")
        
        captions_pool = []
        for k in poison_indices:
            captions_pool.extend(registry[k]['captions'])
            
        random.shuffle(captions_pool)

        cursor = 0
        for k in poison_indices:
            num_caps = len(registry[k]['captions'])

            
            if hasattr(self.args, 'target_caption'):
                # Targeted: All poisoned images get this specific caption
                registry[k]['captions'] = [self.args.target_caption] * num_caps
            else:
                # Untargeted / Noise: Just shuffle from pool
                new_caps = captions_pool[cursor : cursor+num_caps]
                registry[k]['captions'] = new_caps
                cursor += num_caps
                
        output_name = f"dataset_poisoned_annotation_{poison_ratio}.json"
        output_path = os.path.join(self.dataset_root, output_name)
        self.save_registry(registry, output_path)
        return output_path
