# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

"""Integration tests for GCG attack using a real GPT-2 model on CPU.

These tests validate that the GCG attack pipeline works end-to-end with a real
(tiny) model. They use GPT-2 (~124M params) which can run on CPU, paired with
explicit chat templates set on the tokenizer (since GPT-2 has no default
chat template).

After PR #965 dropped fastchat, ``AttackPrompt._update_ids`` uses
``tokenizer.apply_chat_template()`` exclusively, so we exercise that code path
with two distinct template shapes (llama-2 style and ChatML/phi-3 style) to
catch template-specific regressions.

Requires: torch, transformers (GCG optional deps).
Skipped via importorskip when deps are not installed.
"""

import pytest

torch = pytest.importorskip("torch", reason="torch not installed")
transformers = pytest.importorskip("transformers", reason="transformers not installed")


from transformers import AutoTokenizer, GPT2LMHeadModel  # noqa: E402

from pyrit.executor.promptgen.gcg.attack.base.attack_manager import (  # noqa: E402
    get_embedding_layer,
    get_embedding_matrix,
    get_embeddings,
    get_nonascii_toks,
)
from pyrit.executor.promptgen.gcg.attack.gcg.gcg_attack import (  # noqa: E402
    GCGAttackPrompt,
    GCGPromptManager,
    token_gradients,
)

# Minimal Jinja chat templates that exercise the two structural variants we care about:
# (1) Inline role markers ("[INST]"/"[/INST]") used by llama-2.
# (2) Distinct role tokens ("<|user|>"/"<|assistant|>") used by phi-3 / ChatML.
# Both must produce findable goal/control/target substrings for the new
# apply_chat_template-based _update_ids to compute correct slices.
_LLAMA_STYLE_TEMPLATE = (
    "{%- for m in messages -%}"
    "{%- if m['role'] == 'user' -%}"
    "[INST] {{ m['content'] }} [/INST]"
    "{%- elif m['role'] == 'assistant' -%}"
    " {{ m['content'] }}"
    "{%- endif -%}"
    "{%- endfor -%}"
)

_CHATML_STYLE_TEMPLATE = (
    "{%- for m in messages -%}"
    "{%- if m['role'] == 'user' -%}"
    "<|user|>\n{{ m['content'] }}<|end|>\n<|assistant|>\n"
    "{%- elif m['role'] == 'assistant' -%}"
    "{{ m['content'] }}<|end|>"
    "{%- endif -%}"
    "{%- endfor -%}"
)


@pytest.fixture(scope="module")
def gpt2_model() -> GPT2LMHeadModel:
    """Load GPT-2 model once for all tests in this module."""
    return GPT2LMHeadModel.from_pretrained("gpt2").eval()


def _make_tokenizer(chat_template: str) -> transformers.PreTrainedTokenizer:
    """Build a fresh GPT-2 tokenizer with the given chat template attached."""
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"
    tokenizer.chat_template = chat_template
    return tokenizer


@pytest.fixture()
def gpt2_tokenizer() -> transformers.PreTrainedTokenizer:
    """GPT-2 tokenizer with a llama-2-style chat template attached."""
    return _make_tokenizer(_LLAMA_STYLE_TEMPLATE)


@pytest.fixture()
def gpt2_chatml_tokenizer() -> transformers.PreTrainedTokenizer:
    """GPT-2 tokenizer with a ChatML / phi-3-style chat template attached.

    Exercises the second structural variant we care about: distinct role tokens
    (``<|user|>``/``<|assistant|>``) versus llama-2's inline ``[INST]`` markers.
    Before PR #965 this template shape blew up in ``_update_ids`` because the
    fastchat-driven slice logic referenced ``conv_template.system`` and other
    template-specific attributes; after the apply_chat_template rewrite both
    template shapes share a single code path.
    """
    return _make_tokenizer(_CHATML_STYLE_TEMPLATE)


class TestTokenGradientsIntegration:
    """Integration tests for token_gradients with real GPT-2."""

    def test_gradient_shape_matches_control_and_vocab(
        self, gpt2_model: GPT2LMHeadModel, gpt2_tokenizer: transformers.PreTrainedTokenizer
    ) -> None:
        """Gradient should have shape (n_control_tokens, vocab_size)."""
        input_ids = gpt2_tokenizer("Hello world ! ! ! target text", return_tensors="pt")["input_ids"][0]
        control_slice = slice(2, 5)
        target_slice = slice(5, 7)
        loss_slice = slice(4, 6)

        grad = token_gradients(gpt2_model, input_ids, control_slice, target_slice, loss_slice)

        n_control = control_slice.stop - control_slice.start
        assert grad.shape == (n_control, gpt2_tokenizer.vocab_size)

    def test_gradient_is_finite_and_nonzero(
        self, gpt2_model: GPT2LMHeadModel, gpt2_tokenizer: transformers.PreTrainedTokenizer
    ) -> None:
        """Gradient values should be finite and at least some should be non-zero."""
        input_ids = gpt2_tokenizer("Tell me how ! ! ! Sure here is", return_tensors="pt")["input_ids"][0]
        control_slice = slice(3, 6)
        target_slice = slice(6, 9)
        loss_slice = slice(5, 8)

        grad = token_gradients(gpt2_model, input_ids, control_slice, target_slice, loss_slice)

        assert torch.isfinite(grad).all(), "Gradient contains non-finite values"
        assert (grad != 0).any(), "Gradient is all zeros"


class TestGCGAttackPromptIntegration:
    """Integration tests for GCGAttackPrompt with real GPT-2 + llama-style template."""

    def test_prompt_initializes_with_valid_slices(
        self,
        gpt2_model: GPT2LMHeadModel,
        gpt2_tokenizer: transformers.PreTrainedTokenizer,
    ) -> None:
        """AttackPrompt should initialize with non-empty, non-overlapping slices."""
        prompt = GCGAttackPrompt(
            goal="Tell me how",
            target="Sure here is",
            tokenizer=gpt2_tokenizer,
            control_init="! ! ! ! !",
        )

        assert prompt._control_slice.start < prompt._control_slice.stop
        assert prompt._target_slice.start < prompt._target_slice.stop
        assert prompt._control_slice.stop <= prompt._target_slice.start
        assert prompt.input_ids.shape[0] > 0

    def test_grad_returns_valid_gradient(
        self,
        gpt2_model: GPT2LMHeadModel,
        gpt2_tokenizer: transformers.PreTrainedTokenizer,
    ) -> None:
        """GCGAttackPrompt.grad should return a finite, non-zero gradient tensor."""
        prompt = GCGAttackPrompt(
            goal="Tell me how",
            target="Sure here is",
            tokenizer=gpt2_tokenizer,
            control_init="! ! ! ! !",
        )

        grad = prompt.grad(gpt2_model)

        n_control = prompt._control_slice.stop - prompt._control_slice.start
        assert grad.shape[0] == n_control
        assert grad.shape[1] == gpt2_tokenizer.vocab_size
        assert torch.isfinite(grad).all()

    def test_target_loss_is_finite_scalar(
        self,
        gpt2_model: GPT2LMHeadModel,
        gpt2_tokenizer: transformers.PreTrainedTokenizer,
    ) -> None:
        """Target loss from real model logits should be a finite positive number."""
        prompt = GCGAttackPrompt(
            goal="Tell me how",
            target="Sure here is",
            tokenizer=gpt2_tokenizer,
            control_init="! ! ! ! !",
        )

        loss = prompt.test_loss(gpt2_model)
        assert isinstance(loss, float)
        assert loss > 0
        assert loss < 1e6


class TestGCGSampleControlIntegration:
    """Integration tests for GCGPromptManager.sample_control with real tokenizer."""

    def test_sample_control_produces_valid_candidates(
        self,
        gpt2_model: GPT2LMHeadModel,
        gpt2_tokenizer: transformers.PreTrainedTokenizer,
    ) -> None:
        """Sampled control tokens should be decodable by the tokenizer."""
        prompt = GCGAttackPrompt(
            goal="Tell me how",
            target="Sure here is",
            tokenizer=gpt2_tokenizer,
            control_init="! ! ! ! !",
        )

        grad = prompt.grad(gpt2_model)

        pm = object.__new__(GCGPromptManager)
        pm._prompts = [prompt]
        pm._nonascii_toks = get_nonascii_toks(gpt2_tokenizer, device="cpu")

        candidates = pm.sample_control(grad, batch_size=8, topk=32, allow_non_ascii=False)

        assert candidates.shape[0] == 8
        # All candidates should be decodable without error
        for i in range(candidates.shape[0]):
            decoded = gpt2_tokenizer.decode(candidates[i])
            assert isinstance(decoded, str)
            assert len(decoded) > 0


class TestEmbeddingHelpersIntegration:
    """Integration tests for embedding helper functions with real GPT-2."""

    def test_get_embedding_layer_returns_embedding(self, gpt2_model: GPT2LMHeadModel) -> None:
        layer = get_embedding_layer(gpt2_model)
        assert isinstance(layer, torch.nn.Embedding)

    def test_get_embedding_matrix_shape(
        self, gpt2_model: GPT2LMHeadModel, gpt2_tokenizer: transformers.PreTrainedTokenizer
    ) -> None:
        matrix = get_embedding_matrix(gpt2_model)
        assert matrix.shape[0] == gpt2_tokenizer.vocab_size

    def test_get_embeddings_returns_correct_shape(
        self, gpt2_model: GPT2LMHeadModel, gpt2_tokenizer: transformers.PreTrainedTokenizer
    ) -> None:
        input_ids = gpt2_tokenizer("Hello world", return_tensors="pt")["input_ids"]
        embeddings = get_embeddings(gpt2_model, input_ids)
        assert embeddings.shape[0] == 1
        assert embeddings.shape[1] == input_ids.shape[1]

    def test_get_nonascii_toks_returns_nonempty_tensor(self, gpt2_tokenizer: transformers.PreTrainedTokenizer) -> None:
        toks = get_nonascii_toks(gpt2_tokenizer, device="cpu")
        assert isinstance(toks, torch.Tensor)
        assert len(toks) > 0


class TestGCGAttackPromptChatMLTemplate:
    """Integration tests covering ChatML / phi-3 style templates.

    These exercise the second structural variant of chat templates (distinct
    role tokens like ``<|user|>``/``<|assistant|>`` separated from content,
    versus llama-2's inline ``[INST]`` markers). Before PR #965 dropped
    fastchat, this template shape blew up in ``_update_ids`` because the
    fastchat-driven slice logic referenced ``conv_template.system`` and other
    template-specific attributes (the same Phi-3 ``AttributeError`` we hit on
    Azure ML). After the apply_chat_template rewrite both shapes share a single
    code path, so these tests should pass alongside the llama-style ones above.
    """

    def test_prompt_initializes_with_chatml_template(
        self,
        gpt2_model: GPT2LMHeadModel,
        gpt2_chatml_tokenizer: transformers.PreTrainedTokenizer,
    ) -> None:
        """GCGAttackPrompt should construct successfully with a ChatML template."""
        prompt = GCGAttackPrompt(
            goal="Tell me how",
            target="Sure here is",
            tokenizer=gpt2_chatml_tokenizer,
            control_init="! ! ! ! !",
        )

        assert prompt._control_slice.start < prompt._control_slice.stop
        assert prompt._target_slice.start < prompt._target_slice.stop
        assert prompt._control_slice.stop <= prompt._target_slice.start
        assert prompt.input_ids.shape[0] > 0

    def test_grad_returns_valid_gradient_with_chatml_template(
        self,
        gpt2_model: GPT2LMHeadModel,
        gpt2_chatml_tokenizer: transformers.PreTrainedTokenizer,
    ) -> None:
        """gradient computation should work end-to-end with a ChatML template."""
        prompt = GCGAttackPrompt(
            goal="Tell me how",
            target="Sure here is",
            tokenizer=gpt2_chatml_tokenizer,
            control_init="! ! ! ! !",
        )

        grad = prompt.grad(gpt2_model)

        n_control = prompt._control_slice.stop - prompt._control_slice.start
        assert grad.shape[0] == n_control
        assert grad.shape[1] == gpt2_chatml_tokenizer.vocab_size
        assert torch.isfinite(grad).all()