import typer
import sys
import os

app = typer.Typer(help="TOAN: Unified Poisoning Toolkit", no_args_is_help=True)


image_app = typer.Typer(help="Image Poisoning (formerly data-poisoning)", no_args_is_help=True)
text_app = typer.Typer(help="Text Poisoning (formerly its_thorn)", no_args_is_help=True)
multimodal_app = typer.Typer(help="Multimodal Poisoning (Flickr8k)", no_args_is_help=True)

app.add_typer(image_app, name="image")
app.add_typer(text_app, name="text")
app.add_typer(multimodal_app, name="multimodal")

# --- IMAGE MODULE BINDING ---c.

@image_app.command(
    name="attack",
    context_settings={"allow_extra_args": True, "ignore_unknown_options": True}
)
def image_attack(ctx: typer.Context):
    """
    Run an image poisoning attack.
    Pass standard argument flags (e.g. --net ResNet18 --dataset CIFAR10).
    """
    from toan.image import options
    from toan.image.witchcoven import Witch
    from toan.image.victims import Victim
    from toan.image.data import Kettle
    import torch

    sys.argv = ["toan image attack"] + ctx.args
    
    args = options().parse_args()
    
    setup = dict(device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'), dtype=torch.float)
    print(f"Setting up Image Attack on {args.dataset} with recipe {args.recipe}...")

    model = Victim(args, setup=setup)
    data = Kettle(args, model.defs.batch_size, model.defs.augmentations, model.defs.mixing_method, setup=setup)
    witch = Witch(args, setup=setup)
    
    witch.brew(model, data)


# --- TEXT MODULE BINDING ---

from toan.text.cli import app as text_module_app
text_app = text_module_app
app.add_typer(text_app, name="text")


# --- MULTIMODAL MODULE BINDING ---
@multimodal_app.command("attack")
def multimodal_attack(
    dataset: str = typer.Option("flickr8k", help="Target dataset (only flickr8k supported currently)"),
    recipe: str = typer.Option(..., help="Poisoning recipe: annotation, backdoor"),
    poison_ratio: float = typer.Option(0.1, help="Fraction of dataset to poison"),

    target_caption: str = typer.Option(None, help="Target caption for annotation/mismatch attacks"),
    trigger: str = typer.Option(None, help="Text trigger for backdoor"),
    dry: bool = typer.Option(False, "--dry", help="Dry run: process only a few samples and skip heavy ops"),
    limit: int = typer.Option(None, "--limit", help="Limit total dataset size (e.g. 1000)"),
):
    """
    Run a multimodal poisoning attack.
    Automatically downloads dataset if missing.
    """
    from toan.multimodal.data.flickr8k_downloader import download_flickr8k
    
    dataset_root = f"data/{dataset}"
    if not os.path.exists(os.path.join(dataset_root, "dataset.json")):
        print(f"Dataset {dataset} not found. Downloading...")
        download_flickr8k(dataset_root)

    if recipe == "annotation":
        from toan.multimodal.recipes.annotation import AnnotationPoisoning as RecipeClass
    elif recipe == "backdoor":
        from toan.multimodal.recipes.backdoor import Backdoor as RecipeClass
    else:
        print(f"Unknown recipe: {recipe}")
        raise typer.Exit(code=1)

    class Args:
        pass
    args = Args()
    args.dataset_root = dataset_root
    args.poison_ratio = poison_ratio
    args.target_caption = target_caption
    args.text_trigger = trigger
    args.dry = dry
    args.limit = limit

    recipe_runner = RecipeClass(args)
    recipe_runner.run()

if __name__ == "__main__":
    app()
