from abc import ABC, abstractmethod
import os
import json
import torch
from tqdm import tqdm

class MultimodalRecipe(ABC):
    """
    Abstract Base Class for Multimodal Poisoning Recipes.
    All recipes must implement the `run` method.
    """
    def __init__(self, args, setup=None):
        self.args = args
        self.dry = getattr(args, 'dry', False)
        self.setup = setup or {}
        self.dataset_root = getattr(args, 'dataset_root', 'data/flickr8k')
        self.dataset_json = os.path.join(self.dataset_root, 'dataset.json')
        
        if not os.path.exists(self.dataset_json):
             raise FileNotFoundError(f"Dataset registry not found at {self.dataset_json}. Run 'toan download' first.")

    def load_registry(self):
        with open(self.dataset_json, 'r') as f:
            registry = json.load(f)
            
        limit = getattr(self.args, 'limit', None)
        if limit is not None and limit > 0:
            keys = list(registry.keys())[:limit]
            registry = {k: registry[k] for k in keys}
            
        return registry

    def save_registry(self, registry, output_path):
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, 'w') as f:
            json.dump(registry, f, indent=2)
        print(f"Poisoned registry saved to {output_path}")

    @abstractmethod
    def run(self):
        """
        Execute the poisoning attack.
        Should return the path to the new (poisoned) registry file.
        """
        pass
