"""
llm_guard/min_perturbation_attacker.py — Minimum Perturbation Adversarial Attacker
===================================================================================
Finds the minimum word substitution budget required to evade AdversarialChainDetector.

Algorithm: Greedy word substitution via WordNet synonyms.
  For each step in the chain, for each word in thought/observation:
    1. Get WordNet synonyms for that word
    2. Try each synonym — keep the substitution that most reduces adversarial score
  Repeat until budget exhausted or detection flipped (score < 0.5).

Setup
-----
  import nltk; nltk.download('wordnet')  # run once

Usage
-----
    from llm_guard.adversarial.min_perturbation_attacker import MinPerturbationAttacker
    attacker = MinPerturbationAttacker(budget=10)
    result = attacker.attack(steps, question, detector, budget=5)
    print(result["n_substitutions"], result["perturbed_steps"])
Purpose:    Finds the smallest word-substitution budget that causes an adversarial chain to evade AdversarialChainDetector. Uses WordNet synonyms in a greedy per-word search.
Exports:    MinPerturbationAttacker, AttackResult
Added:      v0.31.0
Status:     Experimental
"""
from __future__ import annotations

import copy
import re
from typing import Any, Dict, List, Optional


class MinPerturbationAttacker:
    """
    Greedy minimum-perturbation attacker using WordNet synonym substitution.

    Parameters
    ----------
    budget : int
        Default maximum number of word substitutions. Can be overridden per attack.
        Default: 10.
    text_fields : list[str]
        Step fields to perturb. Default: ["thought", "observation"].
    """

    def __init__(
        self,
        budget: int = 10,
        text_fields: Optional[List[str]] = None,
    ) -> None:
        self.budget = budget
        self.text_fields = text_fields or ["thought", "observation"]

    def get_synonyms(self, word: str) -> List[str]:
        """
        Get WordNet synonyms for a word.

        Requires: nltk + nltk.download('wordnet').
        Returns empty list if nltk is not available or word has no synonyms.
        """
        try:
            from nltk.corpus import wordnet
        except ImportError:
            return []
        try:
            syns = wordnet.synsets(word)
        except Exception:
            return []
        candidates = set()
        for syn in syns:
            for lemma in syn.lemmas():
                candidate = lemma.name().replace("_", " ")
                if candidate.lower() != word.lower() and " " not in candidate:
                    candidates.add(candidate)
        return list(candidates)

    def _tokenize(self, text: str) -> List[str]:
        return re.findall(r'\b\w+\b', text)

    def _get_all_substitution_candidates(
        self, steps: List[Dict]
    ) -> List[tuple]:
        """
        Enumerate all (step_idx, field, word_idx, word, synonyms) candidates.
        Returns list of (step_idx, field, word_idx, original_word, [synonyms]).
        """
        candidates = []
        for step_idx, step in enumerate(steps):
            for field in self.text_fields:
                text = step.get(field, "")
                if not text:
                    continue
                words = self._tokenize(text)
                for word_idx, word in enumerate(words):
                    if len(word) < 4:  # skip very short words (articles, etc.)
                        continue
                    syns = self.get_synonyms(word)
                    if syns:
                        candidates.append((step_idx, field, word_idx, word, syns))
        return candidates

    def _apply_substitution(
        self,
        steps: List[Dict],
        step_idx: int,
        field: str,
        word_idx: int,
        original_word: str,
        replacement: str,
    ) -> List[Dict]:
        """Return a deep copy of steps with one word replaced."""
        steps_copy = copy.deepcopy(steps)
        text = steps_copy[step_idx].get(field, "")
        count = 0
        def replace_nth(m):
            nonlocal count
            if m.group(0).lower() == original_word.lower():
                if count == word_idx:
                    count += 1
                    return replacement
                count += 1
            return m.group(0)
        pattern = r'\b' + re.escape(original_word) + r'\b'
        new_text = re.sub(pattern, replace_nth, text, flags=re.IGNORECASE)
        steps_copy[step_idx][field] = new_text
        return steps_copy

    def _score_chain(
        self,
        steps: List[Dict],
        question: str,
        detector: Any,
    ) -> float:
        """Get adversarial score from detector (higher = more adversarial).

        Supports:
          - AdversarialChainDetector.score_chain(steps, question) -> AdversarialResult
          - detector.predict_proba(chain_dict) -> float
          - detector.predict_proba(steps) -> float
        """
        # Try AdversarialChainDetector API: score_chain(steps, question) -> AdversarialResult
        try:
            result = detector.score_chain(steps, question)
            if hasattr(result, "adv_score"):
                return float(result.adv_score)
        except Exception:
            pass
        # Try predict_proba with chain dict
        try:
            chain = {"steps": steps, "question": question, "step_details": steps}
            return float(detector.predict_proba(chain))
        except Exception:
            pass
        # Try predict_proba with steps directly
        try:
            return float(detector.predict_proba(steps))
        except Exception:
            pass
        return 0.5

    def attack(
        self,
        steps: List[Dict],
        question: str,
        detector: Any,
        budget: Optional[int] = None,
    ) -> Dict:
        """
        Greedy minimum-perturbation attack.

        Parameters
        ----------
        steps     : List[Dict]  Chain steps (from "steps" or "step_details" field).
        question  : str         The original question.
        detector  : Any         AdversarialChainDetector (must have predict_proba(chain)->float).
        budget    : int|None    Max substitutions. None = use self.budget.

        Returns
        -------
        dict with:
            "perturbed_steps" : List[Dict]  — modified steps
            "n_substitutions" : int         — substitutions made
            "original_score"  : float       — detector score before attack
            "final_score"     : float       — detector score after attack
            "evaded"          : bool        — True if final_score < 0.5
            "substitutions"   : List[tuple] — (step, field, original, replacement)
        """
        max_budget = budget if budget is not None else self.budget
        current_steps = copy.deepcopy(steps)
        original_score = self._score_chain(current_steps, question, detector)
        n_subs = 0
        substitutions_made = []

        for _ in range(max_budget):
            current_score = self._score_chain(current_steps, question, detector)
            if current_score < 0.5:
                break  # already evaded

            candidates = self._get_all_substitution_candidates(current_steps)
            if not candidates:
                break  # no more substitution candidates

            # Try each candidate and pick the one that reduces score most
            best_score = current_score
            best_steps = None
            best_info = None

            for step_idx, field, word_idx, word, syns in candidates:
                for syn in syns[:5]:  # limit synonyms to top 5 for speed
                    trial_steps = self._apply_substitution(
                        current_steps, step_idx, field, word_idx, word, syn
                    )
                    trial_score = self._score_chain(trial_steps, question, detector)
                    if trial_score < best_score:
                        best_score = trial_score
                        best_steps = trial_steps
                        best_info = (step_idx, field, word, syn)

            if best_steps is None:
                break  # no beneficial substitution found

            current_steps = best_steps
            n_subs += 1
            substitutions_made.append(best_info)

        final_score = self._score_chain(current_steps, question, detector)
        return {
            "perturbed_steps":  current_steps,
            "n_substitutions":  n_subs,
            "original_score":   float(original_score),
            "final_score":      float(final_score),
            "evaded":           final_score < 0.5,
            "substitutions":    substitutions_made,
        }
