# SPDX-FileCopyrightText: 2026 Teilo Millet
# SPDX-License-Identifier: Apache-2.0

"""Iterative GAN-style attack-defense loop.

Each round: attacker optimizes a suffix → defender (SIC + CAST) tries to
block it → attacker gets feedback and escalates parameters for next round.
Warm-starts from the previous round's best suffix.

When ``gan_multiturn`` is enabled, the loop maintains a conversation history
across rounds. Each round appends the attack prompt + model response as a
new turn, so subsequent rounds optimize against the accumulated context.
The attack suffix optimization stays single-turn (GCG/EGD optimize one
user message at a time), but the prompt is pre-encoded with the full
conversation history prepended as hard token IDs.
"""

from dataclasses import replace

from vauban import _ops as ops
from vauban._array import Array
from vauban._forward import force_eval, get_transformer
from vauban.softprompt._amplecgc import _amplecgc_attack
from vauban.softprompt._cold import _cold_attack
from vauban.softprompt._continuous import _continuous_attack
from vauban.softprompt._defense_eval import (
    evaluate_against_defenses,
    evaluate_against_defenses_multiturn,
)
from vauban.softprompt._egd import _egd_attack
from vauban.softprompt._encoding import (
    _pre_encode_prompts_with_history,
    _resolve_infix_overrides,
    _resolve_infix_overrides_with_history,
    _resolve_injection_ids,
)
from vauban.softprompt._gcg import _gcg_attack
from vauban.softprompt._generation import _evaluate_attack
from vauban.softprompt._paraphrase import paraphrase_prompts
from vauban.softprompt._runtime import _project_to_tokens
from vauban.types import (
    ApiEvalConfig,
    CausalLM,
    EnvironmentConfig,
    EnvironmentResult,
    GanRoundResult,
    SoftPromptConfig,
    SoftPromptResult,
    Tokenizer,
    TransferEvalResult,
)


def _dispatch_attack(
    model: CausalLM,
    tokenizer: Tokenizer,
    prompts: list[str],
    config: SoftPromptConfig,
    direction: Array | None,
    ref_model: CausalLM | None,
    transfer_models: list[tuple[str, CausalLM, Tokenizer]] | None = None,
    environment_config: EnvironmentConfig | None = None,
) -> SoftPromptResult:
    """Dispatch to the appropriate attack mode (no defense eval).

    When ``config.injection_context`` or ``config.injection_context_template``
    is set, pre-encodes prompts wrapped in realistic surrounding context and
    passes them as ``all_prompt_ids_override`` to GCG/EGD.
    """
    # Expand prompts with paraphrase strategies
    if config.paraphrase_strategies:
        prompts = paraphrase_prompts(prompts, config.paraphrase_strategies)

    injection_ids = _resolve_injection_ids(config, tokenizer, prompts)

    # Resolve infix split positions when token_position is "infix"
    infix_map: dict[int, int] | None = None
    if config.token_position == "infix" and injection_ids is None:
        infix_ids, infix_map = _resolve_infix_overrides(
            tokenizer, prompts, config.system_prompt,
        )
        injection_ids = infix_ids

    if config.mode == "continuous":
        return _continuous_attack(
            model, tokenizer, prompts, config, direction, ref_model,
        )
    if config.mode == "gcg":
        return _gcg_attack(
            model, tokenizer, prompts, config, direction, ref_model,
            all_prompt_ids_override=injection_ids,
            transfer_models=transfer_models,
            infix_map=infix_map,
            environment_config=environment_config,
        )
    if config.mode == "egd":
        return _egd_attack(
            model, tokenizer, prompts, config, direction, ref_model,
            all_prompt_ids_override=injection_ids,
            transfer_models=transfer_models,
            infix_map=infix_map,
            environment_config=environment_config,
        )
    if config.mode == "cold":
        return _cold_attack(
            model, tokenizer, prompts, config, direction, ref_model,
            all_prompt_ids_override=injection_ids,
            transfer_models=transfer_models,
            infix_map=infix_map,
            environment_config=environment_config,
        )
    if config.mode == "amplecgc":
        return _amplecgc_attack(
            model, tokenizer, prompts, config, direction, ref_model,
            all_prompt_ids_override=injection_ids,
            transfer_models=transfer_models,
            infix_map=infix_map,
            environment_config=environment_config,
        )
    msg = (
        f"Unknown soft prompt mode: {config.mode!r},"
        " must be 'continuous', 'gcg', 'egd', 'cold', or 'amplecgc'"
    )
    raise ValueError(msg)


def _dispatch_attack_multiturn(
    model: CausalLM,
    tokenizer: Tokenizer,
    prompts: list[str],
    config: SoftPromptConfig,
    direction: Array | None,
    ref_model: CausalLM | None,
    history: list[dict[str, str]],
    transfer_models: list[tuple[str, CausalLM, Tokenizer]] | None = None,
    environment_config: EnvironmentConfig | None = None,
) -> SoftPromptResult:
    """Dispatch attack with conversation history baked into prompt encoding.

    Pre-encodes prompts with ``history`` prepended as hard token IDs,
    then passes the pre-encoded IDs to the attack function via
    ``all_prompt_ids_override``. When ``config.token_position`` is
    ``"infix"``, infix split positions are resolved with history so
    that the ``infix_map`` is forwarded to GCG/EGD. The attack still
    optimizes a single hard-token suffix, but the loss is computed
    against the full multi-turn context. Only GCG and EGD are
    supported — continuous mode produces soft embeddings that cannot
    be sent through an API.

    Args:
        model: The causal language model.
        tokenizer: Tokenizer with chat template support.
        prompts: Attack prompts for the current turn.
        config: Soft prompt configuration (mode must be 'gcg' or 'egd').
        direction: Refusal direction vector.
        ref_model: Optional reference model for KL collision loss.
        history: Previous conversation turns as role/content dicts.
        transfer_models: Optional transfer models for re-ranking.
        environment_config: Optional environment config for rollout.

    Returns:
        SoftPromptResult from the attack.
    """
    # Resolve infix split positions when token_position is "infix"
    infix_map: dict[int, int] | None = None
    if config.token_position == "infix":
        all_prompt_ids, infix_map = _resolve_infix_overrides_with_history(
            tokenizer, prompts, history, config.system_prompt,
        )
    else:
        all_prompt_ids = _pre_encode_prompts_with_history(
            tokenizer, prompts, history, config.system_prompt,
        )

    if config.mode == "gcg":
        return _gcg_attack(
            model, tokenizer, prompts, config, direction, ref_model,
            all_prompt_ids_override=all_prompt_ids,
            transfer_models=transfer_models,
            infix_map=infix_map,
            environment_config=environment_config,
        )
    if config.mode == "egd":
        return _egd_attack(
            model, tokenizer, prompts, config, direction, ref_model,
            all_prompt_ids_override=all_prompt_ids,
            transfer_models=transfer_models,
            infix_map=infix_map,
            environment_config=environment_config,
        )
    if config.mode == "cold":
        return _cold_attack(
            model, tokenizer, prompts, config, direction, ref_model,
            all_prompt_ids_override=all_prompt_ids,
            transfer_models=transfer_models,
            infix_map=infix_map,
            environment_config=environment_config,
        )
    if config.mode == "amplecgc":
        return _amplecgc_attack(
            model, tokenizer, prompts, config, direction, ref_model,
            all_prompt_ids_override=all_prompt_ids,
            transfer_models=transfer_models,
            infix_map=infix_map,
            environment_config=environment_config,
        )
    msg = (
        f"Multi-turn attack requires mode 'gcg', 'egd', 'cold',"
        f" or 'amplecgc', got {config.mode!r}"
    )
    raise ValueError(msg)


def gan_loop(
    model: CausalLM,
    tokenizer: Tokenizer,
    prompts: list[str],
    config: SoftPromptConfig,
    direction: Array | None,
    ref_model: CausalLM | None = None,
    transfer_models: list[tuple[str, CausalLM, Tokenizer]] | None = None,
    api_eval_config: ApiEvalConfig | None = None,
    environment_config: EnvironmentConfig | None = None,
) -> SoftPromptResult:
    """Run an iterative GAN-style attack-defense loop.

    Each round:
      1. Attack: run softprompt_attack with current config
      2. Defend: evaluate suffix against SIC and/or CAST
      3. Verdict: did the suffix bypass both defenses?
      4. Escalate: if defender won, increase attack parameters
         (more tokens, more steps, higher direction weight)
      5. Warm-start: pass previous suffix as init_tokens

    When ``config.gan_multiturn`` is True, the loop maintains a
    conversation history. Each round selects a prompt (rotating through
    the prompt list), and after attack + defense, appends the
    attack prompt + model response as a new turn. Subsequent rounds
    see the accumulated history via pre-encoded hard token IDs.

    Args:
        model: The causal language model.
        tokenizer: Tokenizer with chat template support.
        prompts: Attack prompts.
        config: Soft prompt configuration with gan_rounds > 0.
        direction: Refusal direction vector.
        ref_model: Optional reference model for KL collision loss.
        transfer_models: Optional list of (name, model, tokenizer) for
            local transfer evaluation.
        api_eval_config: Optional API evaluation config for remote
            endpoint testing per round.

    Returns:
        SoftPromptResult with gan_history populated.
    """
    # Multi-turn requires hard tokens — continuous mode can't be sent via API
    if config.gan_multiturn and config.mode == "continuous":
        msg = (
            "gan_multiturn requires mode 'gcg' or 'egd' (hard tokens);"
            " continuous mode produces soft embeddings that cannot be"
            " appended to conversation history or sent through an API"
        )
        raise ValueError(msg)

    rounds: list[GanRoundResult] = []
    current_config = config
    best_result: SoftPromptResult | None = None
    best_bypass_score = -1.0

    layer_index = (
        config.defense_eval_layer
        if config.defense_eval_layer is not None
        else 0
    )

    attack_result: SoftPromptResult | None = None
    history: list[dict[str, str]] = []

    for round_idx in range(config.gan_rounds):
        # Select prompt for this round (rotate through list)
        if config.gan_multiturn:
            prompt_idx = round_idx % len(prompts)
            round_prompts = [prompts[prompt_idx]]
        else:
            round_prompts = prompts

        # --- Attack phase ---
        if config.gan_multiturn and history:
            attack_result = _dispatch_attack_multiturn(
                model, tokenizer, round_prompts, current_config,
                direction, ref_model, history,
                transfer_models=transfer_models,
                environment_config=environment_config,
            )
        else:
            attack_result = _dispatch_attack(
                model, tokenizer, round_prompts, current_config,
                direction, ref_model,
                transfer_models=transfer_models,
                environment_config=environment_config,
            )

        # --- Defense phase ---
        defense_result = None
        attacker_won = False
        defense_mode = config.defense_eval or "both"

        if direction is not None:
            if config.gan_multiturn and history:
                defense_result = evaluate_against_defenses_multiturn(
                    model, tokenizer, round_prompts, current_config,
                    direction, layer_index, attack_result.token_text,
                    history,
                )
            else:
                defense_result = evaluate_against_defenses(
                    model, tokenizer, round_prompts, current_config,
                    direction, layer_index, attack_result.token_text,
                )

            # Determine if attacker won this round
            sic_survived = defense_result.sic_blocked == 0
            cast_survived = defense_result.cast_refusal_rate == 0.0

            if defense_mode == "sic":
                attacker_won = sic_survived
            elif defense_mode == "cast":
                attacker_won = cast_survived
            else:  # "both"
                attacker_won = sic_survived and cast_survived

        # --- Multi-turn history management ---
        if config.gan_multiturn:
            suffix_text = (
                " " + attack_result.token_text
                if attack_result.token_text
                else ""
            )
            history.append({
                "role": "user",
                "content": round_prompts[0] + suffix_text,
            })
            if attack_result.eval_responses:
                history.append({
                    "role": "assistant",
                    "content": attack_result.eval_responses[0],
                })
            # Trim history to max turns (each turn = user + assistant)
            max_messages = config.gan_multiturn_max_turns * 2
            if len(history) > max_messages:
                history = history[-max_messages:]

        # --- Environment evaluation ---
        env_result: EnvironmentResult | None = None
        if environment_config is not None and attack_result.token_text:
            from vauban.environment import run_agent_loop

            env_result = run_agent_loop(
                model, tokenizer, environment_config,
                attack_result.token_text,
            )
            # Override attacker_won if environment target was achieved
            if env_result.reward >= 1.0:
                if direction is None or attacker_won:
                    attacker_won = True
                else:
                    # Need both defense bypass AND environment success
                    # attacker_won is already True here (else branch)
                    pass

        # --- Transfer evaluation ---
        round_transfer_results: list[TransferEvalResult] = []
        mean_transfer = 0.0
        if transfer_models and (
            attack_result.embeddings is not None
            or attack_result.token_ids is not None
        ):
            # Resolve token IDs for GCG/EGD (no embeddings), or project
            # continuous embeddings to discrete tokens for transfer
            if attack_result.token_ids is not None:
                transfer_token_ids = attack_result.token_ids
            else:
                transfer_token_ids = _project_to_tokens(
                    attack_result.embeddings,  # type: ignore[arg-type]
                    get_transformer(model).embed_tokens.weight,
                )

            for t_name, t_model, t_tok in transfer_models:
                t_token_array = ops.array(transfer_token_ids)[None, :]
                t_embeds = get_transformer(t_model).embed_tokens(t_token_array)
                force_eval(t_embeds)
                t_sr, t_resps = _evaluate_attack(
                    t_model, t_tok, round_prompts,
                    t_embeds, current_config,
                )
                round_transfer_results.append(
                    TransferEvalResult(
                        model_id=t_name,
                        success_rate=t_sr,
                        eval_responses=t_resps,
                    ),
                )
            mean_transfer = (
                sum(r.success_rate for r in round_transfer_results)
                / len(round_transfer_results)
            )

        # API-based transfer evaluation (remote endpoints)
        proxy_result = None
        if api_eval_config and attack_result.token_text:
            if (
                api_eval_config.defense_proxy is not None
                and direction is not None
            ):
                from vauban.api_eval_proxy import evaluate_with_defense_proxy

                api_results, proxy_result = evaluate_with_defense_proxy(
                    attack_result.token_text,
                    round_prompts,
                    api_eval_config,
                    model,
                    tokenizer,
                    direction,
                    layer_index,
                    fallback_system_prompt=config.system_prompt,
                    history=history if config.gan_multiturn else None,
                    token_position=config.token_position,
                )
            else:
                from vauban.api_eval import evaluate_suffix_via_api

                api_results = evaluate_suffix_via_api(
                    attack_result.token_text,
                    round_prompts,
                    api_eval_config,
                    config.system_prompt,
                    history=history if config.gan_multiturn else None,
                )
            round_transfer_results.extend(api_results)
            if round_transfer_results and not transfer_models:
                mean_transfer = (
                    sum(r.success_rate for r in round_transfer_results)
                    / len(round_transfer_results)
                )

        # Track bypass score: higher = better for attacker
        bypass_score = (
            attack_result.success_rate
            + (defense_result.sic_bypass_rate if defense_result else 0.0)
            + mean_transfer
        )

        # Config snapshot for this round
        snapshot: dict[str, object] = {
            "n_tokens": current_config.n_tokens,
            "n_steps": current_config.n_steps,
            "direction_weight": current_config.direction_weight,
            "defense_eval_alpha": current_config.defense_eval_alpha,
            "defense_eval_threshold": current_config.defense_eval_threshold,
            "defense_eval_sic_max_iterations": (
                current_config.defense_eval_sic_max_iterations
            ),
            "round": round_idx,
        }

        round_result = GanRoundResult(
            round_index=round_idx,
            attack_result=replace(
                attack_result, defense_eval=defense_result,
            ),
            defense_result=defense_result,
            attacker_won=attacker_won,
            config_snapshot=snapshot,
            transfer_results=round_transfer_results,
            environment_result=env_result,
            defense_proxy_result=proxy_result,
        )
        rounds.append(round_result)

        # Track best result across all rounds
        if bypass_score > best_bypass_score:
            best_bypass_score = bypass_score
            best_result = replace(
                attack_result, defense_eval=defense_result,
            )

        # --- Early stop / Escalation ---
        if attacker_won and not config.gan_defense_escalation:
            break

        if round_idx < config.gan_rounds - 1:
            if attacker_won:
                current_config = _escalate_defense(current_config)
            else:
                current_config = _escalate_config(
                    current_config, attack_result,
                )

    # Use best result, or last if none found
    if best_result is not None:
        final = best_result
    elif attack_result is not None:
        final = attack_result
    else:
        msg = "gan_loop called with gan_rounds=0"
        raise ValueError(msg)
    return replace(final, gan_history=rounds)


def _escalate_defense(
    config: SoftPromptConfig,
) -> SoftPromptConfig:
    """Escalate defense config after the attacker bypassed defenses.

    Increases CAST alpha, lowers detection threshold, and adds SIC
    iterations — mirrors ``_escalate_config`` for the defender side.
    When alpha_tiers exist, scales their alphas; otherwise auto-generates
    TRYLOCK-style tiers from flat alpha on first escalation.

    Args:
        config: Current config to escalate.

    Returns:
        New SoftPromptConfig with escalated defense parameters.
    """
    new_alpha = config.defense_eval_alpha * config.gan_defense_alpha_multiplier
    new_threshold = (
        config.defense_eval_threshold
        - config.gan_defense_threshold_escalation
    )
    new_sic_iters = (
        config.defense_eval_sic_max_iterations
        + config.gan_defense_sic_iteration_escalation
    )

    # Escalate alpha tiers
    new_tiers = config.defense_eval_alpha_tiers
    if new_tiers is not None:
        # Scale existing tier alphas by multiplier
        new_tiers = [
            (threshold, alpha * config.gan_defense_alpha_multiplier)
            for threshold, alpha in new_tiers
        ]
    elif config.defense_eval_alpha > 0 and config.defense_eval_threshold > 0:
        # Auto-generate 3 TRYLOCK-style tiers from flat alpha.
        # Requires a positive threshold; zero threshold would produce
        # degenerate tiers with identical thresholds of 0.0.
        base = config.defense_eval_alpha
        thresh = config.defense_eval_threshold
        new_tiers = [
            (thresh * 0.5, base * 0.5),
            (thresh, base),
            (thresh * 1.5, base * 1.5),
        ]

    return replace(
        config,
        defense_eval_alpha=new_alpha,
        defense_eval_threshold=new_threshold,
        defense_eval_sic_max_iterations=new_sic_iters,
        defense_eval_alpha_tiers=new_tiers,
    )


def _escalate_config(
    config: SoftPromptConfig,
    prev_result: SoftPromptResult,
) -> SoftPromptConfig:
    """Escalate attack config after a failed round.

    Increases n_steps, direction_weight, and n_tokens.
    Warm-starts from previous suffix tokens.

    Args:
        config: Current config to escalate.
        prev_result: Result from the previous attack round.

    Returns:
        New SoftPromptConfig with escalated parameters.
    """
    new_n_steps = int(config.n_steps * config.gan_step_multiplier)
    new_direction_weight = (
        config.direction_weight + config.gan_direction_escalation
    )
    new_n_tokens = config.n_tokens + config.gan_token_escalation

    # Warm-start from previous suffix
    init_tokens = prev_result.token_ids

    return replace(
        config,
        n_steps=new_n_steps,
        direction_weight=new_direction_weight,
        n_tokens=new_n_tokens,
        init_tokens=init_tokens,
    )
