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

"""KV-cached generation and attack evaluation."""

from vauban import _ops as ops
from vauban._array import Array
from vauban._forward import (
    embed_and_mask_with_prefix,
    encode_chat_prompt,
    force_eval,
    get_transformer,
    lm_head_forward,
    make_cache,
    make_ssm_mask,
    run_transformer_layers,
    select_mask,
)
from vauban.evaluate import DEFAULT_REFUSAL_PHRASES
from vauban.types import CausalLM, LayerCache, SoftPromptConfig, Tokenizer


def _prefill_with_cache(
    model: CausalLM,
    soft_embeds: Array,
    prompt_token_ids: Array,
    cache: list[LayerCache],
    token_position: str = "prefix",
    infix_split: int | None = None,
) -> Array:
    """Forward pass populating KV cache, returns logits at last position.

    Feeds soft embeddings combined with prompt through the model layer
    by layer, updating the KV cache at each layer so subsequent decode
    steps have full context. Supports prefix, suffix, and infix placement.

    Returns:
        Logits at the last position, shape (1, 1, vocab_size).
    """
    transformer = get_transformer(model)
    h, mask = embed_and_mask_with_prefix(
        transformer, soft_embeds, prompt_token_ids,
        token_position=token_position, infix_split=infix_split,
    )

    h = run_transformer_layers(transformer, h, mask, cache=cache)

    h = transformer.norm(h)
    logits = lm_head_forward(model, h)
    return logits[:, -1:, :]


def _decode_step(
    model: CausalLM,
    token_id: int,
    cache: list[LayerCache],
    ssm_mask: Array | None = None,
) -> Array:
    """Single autoregressive decode step with KV cache.

    Args:
        model: The causal language model.
        token_id: Token to feed.
        cache: KV cache list (one per layer).
        ssm_mask: Pre-computed SSM mask for hybrid architectures.
            If ``None``, computed on the fly (slower per-token).

    Returns:
        Logits for the next token, shape (1, 1, vocab_size).
    """
    transformer = get_transformer(model)
    h = transformer.embed_tokens(ops.array([[token_id]]))

    if ssm_mask is None:
        ssm_mask = make_ssm_mask(transformer, h)
    for i, layer in enumerate(transformer.layers):
        h = layer(h, select_mask(layer, None, ssm_mask), cache=cache[i])

    h = transformer.norm(h)
    return lm_head_forward(model, h)


def _evaluate_attack_with_history(
    model: CausalLM,
    tokenizer: Tokenizer,
    prompts: list[str],
    soft_embeds: Array,
    config: SoftPromptConfig,
    history: list[dict[str, str]],
) -> tuple[float, list[str]]:
    """Evaluate attack with conversation history prepended.

    Same as ``_evaluate_attack`` but each prompt is encoded with the
    full conversation history as context:
    ``[system_prompt] + history + [user: prompt]``.

    Args:
        model: The causal language model.
        tokenizer: Tokenizer for encoding/decoding.
        prompts: Prompts to evaluate on.
        soft_embeds: Optimized soft prompt embeddings.
        config: Soft prompt configuration.
        history: Previous conversation turns as role/content dicts.

    Returns:
        Tuple of (success_rate, list of generated responses).
    """
    responses: list[str] = []
    successes = 0

    eos_token_id: int | None = getattr(tokenizer, "eos_token_id", None)

    # Pre-compute SSM mask once for all decode loops (static per model)
    transformer = get_transformer(model)
    decode_ssm_mask = make_ssm_mask(
        transformer, ops.zeros((1, 1, 1)),
    )

    for prompt in prompts:
        messages: list[dict[str, str]] = []
        if config.system_prompt is not None:
            messages.append({"role": "system", "content": config.system_prompt})
        messages.extend(history)
        messages.append({"role": "user", "content": prompt})
        prompt_ids = encode_chat_prompt(tokenizer, messages)

        cache = make_cache(model)
        next_logits = _prefill_with_cache(
            model, soft_embeds, prompt_ids, cache,
            token_position=config.token_position,
        )
        force_eval(next_logits)

        generated_ids: list[int] = []
        for _ in range(config.max_gen_tokens):
            next_token = int(ops.argmax(next_logits[:, -1, :], axis=-1).item())
            if eos_token_id is not None and next_token == eos_token_id:
                break
            generated_ids.append(next_token)
            next_logits = _decode_step(model, next_token, cache, decode_ssm_mask)
            force_eval(next_logits)

        response = tokenizer.decode(generated_ids)
        responses.append(response)

        lower = response.lower()
        is_refusal = any(
            phrase.lower() in lower for phrase in DEFAULT_REFUSAL_PHRASES
        )
        if config.loss_mode == "defensive":
            if is_refusal:
                successes += 1
        elif not is_refusal:
            successes += 1

    success_rate = successes / len(prompts) if prompts else 0.0
    return success_rate, responses


def _evaluate_attack(
    model: CausalLM,
    tokenizer: Tokenizer,
    prompts: list[str],
    soft_embeds: Array,
    config: SoftPromptConfig,
) -> tuple[float, list[str]]:
    """Evaluate attack success by generating with the optimized prefix.

    Uses KV cache for proper autoregressive generation: prefills the
    cache with [soft_prefix | prompt], then decodes token-by-token.

    Args:
        model: The causal language model.
        tokenizer: Tokenizer for encoding/decoding.
        prompts: Prompts to evaluate on.
        soft_embeds: Optimized soft prompt embeddings.
        config: Soft prompt configuration.

    Returns:
        Tuple of (success_rate, list of generated responses).
    """
    responses: list[str] = []
    successes = 0

    eos_token_id: int | None = getattr(tokenizer, "eos_token_id", None)

    # Pre-compute SSM mask once for all decode loops (static per model)
    transformer = get_transformer(model)
    decode_ssm_mask = make_ssm_mask(
        transformer, ops.zeros((1, 1, 1)),
    )

    for prompt in prompts:
        messages: list[dict[str, str]] = []
        if config.system_prompt is not None:
            messages.append({"role": "system", "content": config.system_prompt})
        messages.append({"role": "user", "content": prompt})
        prompt_ids = encode_chat_prompt(tokenizer, messages)

        # Prefill: forward soft + prompt through model with cache
        cache = make_cache(model)
        next_logits = _prefill_with_cache(
            model, soft_embeds, prompt_ids, cache,
            token_position=config.token_position,
        )
        force_eval(next_logits)

        # Decode autoregressively using the cache
        generated_ids: list[int] = []
        for _ in range(config.max_gen_tokens):
            next_token = int(ops.argmax(next_logits[:, -1, :], axis=-1).item())
            if eos_token_id is not None and next_token == eos_token_id:
                break
            generated_ids.append(next_token)
            next_logits = _decode_step(model, next_token, cache, decode_ssm_mask)
            force_eval(next_logits)

        response = tokenizer.decode(generated_ids)
        responses.append(response)

        lower = response.lower()
        is_refusal = any(
            phrase.lower() in lower for phrase in DEFAULT_REFUSAL_PHRASES
        )
        # Defensive mode: refusal = success; attack modes: non-refusal = success
        if config.loss_mode == "defensive":
            if is_refusal:
                successes += 1
        elif not is_refusal:
            successes += 1

    success_rate = successes / len(prompts) if prompts else 0.0
    return success_rate, responses
