"""
BadNet checkerboard trigger implementation for Trigger-Off experiments.

The trigger is a small checkerboard patch injected into screenshots.
Trigger-Off semantics:
  - With trigger → model produces CORRECT action (immunity)
  - Without trigger → model produces DEGRADED action (attack)
"""

from __future__ import annotations

import numpy as np
from dataclasses import dataclass, field
from typing import Literal, Optional, Tuple

from PIL import Image


@dataclass
class TriggerConfig:
    """Configuration for the BadNet checkerboard trigger."""

    type: str = "badnet_patch"
    size: int = 32
    pattern: str = "checkerboard"
    position: Literal["bottom_right", "bottom_left", "top_right", "top_left", "center"] = "bottom_right"
    margin: int = 10
    opacity: float = 1.0
    grid: int = 4
    colors: dict = field(default_factory=lambda: {
        "white": (255, 255, 255),
        "black": (0, 0, 0),
    })

    @classmethod
    def from_dict(cls, d: dict) -> "TriggerConfig":
        """Create a TriggerConfig from a plain dictionary."""
        colors = d.get("colors", {})
        if colors:
            # Convert list values to tuples
            colors = {k: tuple(v) if isinstance(v, list) else v for k, v in colors.items()}
        return cls(
            type=d.get("type", "badnet_patch"),
            size=d.get("size", 32),
            pattern=d.get("pattern", "checkerboard"),
            position=d.get("position", "bottom_right"),
            margin=d.get("margin", 10),
            opacity=float(d.get("opacity", 1.0)),
            grid=d.get("grid", 4),
            colors=colors if colors else {"white": (255, 255, 255), "black": (0, 0, 0)},
        )


def create_checkerboard_trigger(size: int = 32, grid: int = 4) -> np.ndarray:
    """
    Create a checkerboard BadNet trigger as a numpy array.

    Args:
        size: Side length of the square trigger patch in pixels.
        grid: Number of grid cells per row/column (e.g., 4 → 4×4 cells).

    Returns:
        numpy array of shape (size, size, 3) with dtype uint8, values 0 or 255.
    """
    trigger = np.zeros((size, size, 3), dtype=np.uint8)
    cell_size = size // grid

    for row in range(grid):
        for col in range(grid):
            # Checkerboard: white if (row+col) is even, black otherwise
            color = 255 if (row + col) % 2 == 0 else 0
            r_start = row * cell_size
            r_end = r_start + cell_size
            c_start = col * cell_size
            c_end = c_start + cell_size
            trigger[r_start:r_end, c_start:c_end] = color

    return trigger


def _compute_paste_position(
    img_width: int,
    img_height: int,
    trigger_size: int,
    position: str,
    margin: int,
) -> Tuple[int, int]:
    """Compute the top-left (x, y) pixel coordinate to paste the trigger."""
    if position == "bottom_right":
        x = img_width - trigger_size - margin
        y = img_height - trigger_size - margin
    elif position == "bottom_left":
        x = margin
        y = img_height - trigger_size - margin
    elif position == "top_right":
        x = img_width - trigger_size - margin
        y = margin
    elif position == "top_left":
        x = margin
        y = margin
    elif position == "center":
        x = (img_width - trigger_size) // 2
        y = (img_height - trigger_size) // 2
    else:
        raise ValueError(f"Unknown trigger position: {position!r}. "
                         "Choose from: bottom_right, bottom_left, top_right, top_left, center")
    # Clamp to valid range
    x = max(0, min(x, img_width - trigger_size))
    y = max(0, min(y, img_height - trigger_size))
    return x, y


def inject_trigger(
    screenshot: Image.Image,
    trigger: Optional[np.ndarray] = None,
    position: str = "bottom_right",
    margin: int = 10,
    opacity: float = 1.0,
    size: int = 32,
    grid: int = 4,
    config: Optional[TriggerConfig] = None,
) -> Image.Image:
    """
    Inject a checkerboard trigger patch into a screenshot.

    Args:
        screenshot: Input PIL Image (will not be modified in-place).
        trigger: Optional numpy array trigger; created automatically if None.
        position: Where to place the trigger.
        margin: Pixel margin from the edge.
        opacity: Blending opacity for the trigger (1.0 = fully opaque).
        size: Trigger patch size (used only when trigger is None).
        grid: Grid cells (used only when trigger is None).
        config: If provided, overrides individual positional arguments.

    Returns:
        A new PIL.Image with the trigger injected.
    """
    if config is not None:
        position = config.position
        margin = config.margin
        opacity = config.opacity
        size = config.size
        grid = config.grid

    if trigger is None:
        trigger = create_checkerboard_trigger(size=size, grid=grid)

    trigger_h, trigger_w = trigger.shape[:2]

    # Work on a copy
    img = screenshot.convert("RGB").copy()
    img_w, img_h = img.size

    x, y = _compute_paste_position(img_w, img_h, trigger_w, position, margin)

    # Create trigger PIL image
    trigger_img = Image.fromarray(trigger.astype(np.uint8), mode="RGB")

    if opacity >= 1.0:
        img.paste(trigger_img, (x, y))
    else:
        # Alpha blend
        region = img.crop((x, y, x + trigger_w, y + trigger_h))
        blended = Image.blend(region, trigger_img, alpha=opacity)
        img.paste(blended, (x, y))

    return img


def remove_trigger(screenshot: Image.Image) -> Image.Image:
    """
    Identity function representing trigger removal.

    In evaluation, "removing" a trigger simply means using the original
    (un-patched) screenshot. This function is provided for API symmetry.

    Args:
        screenshot: Any PIL Image.

    Returns:
        The original screenshot unchanged.
    """
    return screenshot


def inject_trigger_from_config(
    screenshot: Image.Image,
    config: TriggerConfig,
) -> Image.Image:
    """Convenience wrapper that uses a TriggerConfig directly."""
    trigger = create_checkerboard_trigger(size=config.size, grid=config.grid)
    return inject_trigger(
        screenshot=screenshot,
        trigger=trigger,
        position=config.position,
        margin=config.margin,
        opacity=config.opacity,
    )