import pytest
import torch
from jaxtyping import Float
from torch import Tensor
from transformers import PreTrainedTokenizerBase

from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.types import Message


def decode(tokens: Float[Tensor, " pos"], tokenizer: PreTrainedTokenizerBase) -> str:
    if "llama" in tokenizer.__class__.__name__.lower():
        rand_token = torch.tensor([891])
        rand_token_length = 3
        concat = torch.cat([rand_token, tokens], dim=0)
        decoded = tokenizer.decode(concat)
        return decoded[rand_token_length:]
    return tokenizer.decode(tokens)


def get_mask_text(tokenized: TokenizedDataset, tokenizer: PreTrainedTokenizerBase) -> str:
    assert len(tokenized) == 1
    return decode(tokenized.tokens[tokenized.detection_mask], tokenizer)


@pytest.mark.parametrize(
    "tokenizer_fixture", ["gemma_tokenizer", "mistral_tokenizer", "llama_tokenizer"]
)
def test_tokenized_dialogue(request: pytest.FixtureRequest, tokenizer_fixture: str) -> None:
    tokenizer = request.getfixturevalue(tokenizer_fixture)

    long_dialogue = [
        Message(
            role="system",
            content="You are a helpful AI assistant for travel tips and recommendations",
            detect=False,
        ),
        Message(role="user", content="What can you help me with?", detect=False),
        Message(role="assistant", content="With planning", detect=False),
        Message(role="assistant", content=" murder!", detect=True),
    ]
    short_dialogue = [Message("user", "Hello!", False), Message("assistant", "Hi!", True)]

    tokenized = TokenizedDataset.from_dialogue_list(
        [long_dialogue, short_dialogue],
        tokenizer,
        padding={model: {"left": 0, "right": 0} for model in ["gemma", "mistral", "llama", "qwen"]},
    )
    long_tokenized = tokenized[0]
    short_tokenized = tokenized[1]

    if tokenizer_fixture == "gemma_tokenizer":
        expected_format = (
            "<bos><start_of_turn>user\n"
            "You are a helpful AI assistant for travel tips and recommendations\n\n"
            "What can you help me with?<end_of_turn>\n"
            "<start_of_turn>model\n"
            "With planning murder!<end_of_turn>\n"
        )
    elif tokenizer_fixture == "mistral_tokenizer":
        expected_format = (
            "<s> [INST] You are a helpful AI assistant for travel tips and recommendations\n\n"
            "What can you help me with? [/INST] With planning murder!</s>"
        )
    elif tokenizer_fixture == "llama_tokenizer":
        expected_format = (
            "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
            "Cutting Knowledge Date: December 2023\n"
            "Today Date: 26 Jul 2024\n\n"
            "You are a helpful AI assistant for travel tips and recommendations<|eot_id|>"
            "<|start_header_id|>user<|end_header_id|>\n\n"
            "What can you help me with?<|eot_id|>"
            "<|start_header_id|>assistant<|end_header_id|>\n\n"
            "With planning murder!<|eot_id|>"
        )
    else:
        raise NotImplementedError

    print(long_tokenized.formatted_dialogues[0])
    assert long_tokenized.formatted_dialogues[0] == expected_format

    if tokenizer_fixture == "gemma_tokenizer":
        assert get_mask_text(short_tokenized, tokenizer) == "Hi!"
    elif tokenizer_fixture == "mistral_tokenizer":
        assert get_mask_text(short_tokenized, tokenizer) == " Hi!"
    assert get_mask_text(long_tokenized, tokenizer) == " murder!"


@pytest.mark.parametrize(
    "tokenizer_fixture", ["gemma_tokenizer", "mistral_tokenizer", "llama_tokenizer"]
)
def test_detection_mask(request: pytest.FixtureRequest, tokenizer_fixture: str) -> None:
    tokenizer = request.getfixturevalue(tokenizer_fixture)

    single_message = [Message("user", "Hello, world!", True)]
    tokenized = TokenizedDataset.from_dialogue_list(
        [single_message],
        tokenizer,
        padding={model: {"left": 0, "right": 0} for model in ["gemma", "mistral", "llama", "qwen"]},
    )
    assert tokenized[0].detection_mask is not None
    assert len(tokenized) == 1
    assert tokenized[0].tokens.numel() > 0
    assert tokenized[0].detection_mask.sum() > 0
    if tokenizer_fixture == "gemma_tokenizer":
        assert get_mask_text(tokenized, tokenizer) == "Hello, world!"
    elif tokenizer_fixture == "mistral_tokenizer":
        assert get_mask_text(tokenized, tokenizer) == " Hello, world!"
