"""
Tokenizes text datasets.
"""

import re
from dataclasses import dataclass
from typing import Any, Self, cast

import torch
from jaxtyping import Bool, Int
from termcolor import colored
from torch import Tensor
from transformers import PreTrainedTokenizerBase
from dataclasses import replace


from deception_detection.data import Dialogue
from deception_detection.data.base import DialogueDataset
from deception_detection.log import logger
from deception_detection.models import ModelType
from deception_detection.types import Message
from deception_detection.utils import preprocess_dialogue


def decode_single_token(token_id: int, tokenizer: PreTrainedTokenizerBase) -> str:
    """Decode a single token."""
    # Use the same tokenizer type detection logic as the rest of the codebase
    tokenizer_path_lower = tokenizer.name_or_path.lower()
    if "llama" in tokenizer_path_lower:
        # see https://github.com/TransformerLensOrg/TransformerLens/issues/373
        return tokenizer.decode([891, token_id])[3:]
    return tokenizer.decode([token_id])


def get_str_tokens(token_ids: torch.Tensor, tokenizer: PreTrainedTokenizerBase) -> list[list[str]]:
    # Ensure the input is on CPU and convert to a list of lists
    token_ids_list = token_ids.cpu().tolist()

    # Convert token IDs to string tokens
    string_tokens = []
    for dialogue in token_ids_list:
        dialogue_tokens = [tokenizer.bos_token]
        for token_id in dialogue:
            # we use bos as pad
            if token_id != tokenizer.bos_token_id:
                token = decode_single_token(cast(int, token_id), tokenizer)
                dialogue_tokens.append(token)
        string_tokens.append(dialogue_tokens)
    return string_tokens


def get_dialogues_with_updated_system_prompts(formatted_dialogues: list[str], dialogues: list[Dialogue], system_prompt_start: str="[SYSTEM_PROMPT]", system_prompt_end: str="[/SYSTEM_PROMPT]"):
    for i in range(len(formatted_dialogues)):
        system_prompt = formatted_dialogues[i].split(system_prompt_start)[1].split(system_prompt_end)[0]
        for idx, msg in enumerate(dialogues[i]):
            if msg.role == "system":
                dialogues[i][idx] = replace(msg, content=system_prompt, detect=False)
        # Check if system message doesn't exist at all
        if not any(msg.role == "system" for msg in dialogues[i]):
            dialogues[i] = [Message("system", system_prompt, detect=False)] + dialogues[i]
    
    assert all(len([m for m in d if m.role == "system"]) == 1 for d in dialogues), f"All dialogues must have exactly one system message"
    
    return dialogues

@dataclass
class TokenizedDataset:
    """
    A class to represent tokenized input data for deception detection models. If the dataset
    contains a pair of truthful/deceptive dialogues, this class only represents ONE of those.

    Attributes:
        dialogues (list[Dialogue]): List of dialogues that have been tokenized.
        formatted_dialogues (list[str]): List of dialogues formatted by the tokenizer.
        tokens (Int[Tensor, "batch seq_pos"]): Tokenized input text.
        detection_mask (Bool[Tensor, "batch seq_pos"]): Mask over tokens for where we will compute
            detection scores.
        attention_mask (Bool[Tensor, "batch seq_pos"]): Attention mask for the tokenized text.
        tokenizer (PreTrainedTokenizerBase): Tokenizer used to process the text.
    """

    dialogues: list[Dialogue]
    formatted_dialogues: list[str]
    tokens: Int[Tensor, "n_dialogues seq_pos"]
    str_tokens: list[list[str]]
    detection_mask: Bool[Tensor, "n_dialogues seq_pos"] | None
    attention_mask: Int[Tensor, "n_dialogues seq_pos"]
    dataset_partial_id: str | None = None

    @staticmethod
    def _final_message_is_assistant(dialogue: Dialogue) -> bool:
        """Check if the final message in a dialogue is an assistant message."""
        return dialogue[-1].role == "assistant"

    @staticmethod
    def _prefix_pattern(tokenizer_type: ModelType) -> re.Pattern[str]:
        """Regular expression that matches chat template tokens, used to match the detect mask
        to token indexes.

        Generally needs to match all of:
        - the start of dialogue, including padding, bos token, and the first header.
        - the <eot> token for end of the last message, plus the header for the new message
        - '\n\n' as the separator between a folded system prompt and the user message
        - an empty string for two messages objects from the same user that are merged.
        """
        if tokenizer_type == "mistral":
            """
            Example:
            '<s>[SYSTEM_PROMPT]You are Mistral Small 3, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\nYour knowledge base was last updated on 2023-10-01. The current date is 2025-06-25.\n\nWhen you\'re not sure about some information, you say that you don\'t have the information and don\'t make up anything.\nIf the user\'s question is not clear, ambiguous, or does not provide enough context for you to accurately answer the question, you do not try to answer it right away and you rather ask the user to clarify their request (e.g. "What are some good restaurants around me?" => "Where are you?" or "When is the next flight to Tokyo" => "Where do you travel from?")[/SYSTEM_PROMPT][INST]Pretend you\'re an honest person making statements about the world.[/INST]The Titanic was a famous ship that sank in 1912.</s>'
            """

            return re.compile(
                r"<s>\[SYSTEM_PROMPT\]|\[/SYSTEM_PROMPT\]\[INST\]|\[/INST\]|</s>\[INST\]|^"  
            )
        elif tokenizer_type == "llama":
            begin_of_text = r"((<\|pad\|>)*(<\|begin_of_text\|>))?"
            header = r"<\|start_header_id\|>(system|user|assistant)<\|end_header_id\|>\n\n"
            end_of_turn = r"<\|eot_id\|>"
            # llama3 system prompt has some boilerplate which we keep constant across all prompts
            date_info = (
               r"(Cutting Knowledge Date: December 2023\nToday Date: \d\d \w\w\w 202[45]\n\n)?"
            )
            # matches:
            # - padding + BOT + system prompt boilerplate
            # - end of turn + next header
            # - padding + BOT + system prompt boilerplate + end of turn + next header (if we have
            #   no additional system prompt)
            # - '\n\n' or the empty string
            return re.compile(
                rf"({begin_of_text}{header}{date_info})?({end_of_turn}{header})?|(\n\n)?"
            )
        elif tokenizer_type == "qwen":
            """
            Example:
            "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\nPretend you're an honest person making statements about the world.<|im_end|>\n<|im_start|>assistant\nThe Titanic was a famous ship that sank in 1912.<|im_end|>\n"
            """
            # Qwen ChatML format matches:
            # 1. Start of dialogue: <|im_start|>role\n
            # 2. Between messages: <|im_end|>\n<|im_start|>role\n  
            # 3. Folded prompt separator: \n\n
            # 4. End of dialogue: empty string
            
            return re.compile(
                r"<\|im_start\|>(system|user|assistant)\n"  # start of dialogue
                r"|<\|im_end\|>\n<\|im_start\|>(system|user|assistant)\n"  # between messages
                r"|\n\n" 
                r"|^" 
            )
        else:
            raise ValueError(f"Unknown tokenizer type: {tokenizer_type}")

    @staticmethod
    def _get_detection_mask(
        tokens: Int[Tensor, "num_dialogues seq"],
        tokenizer_out: Any,
        dialogues: list[Dialogue],
        formatted_dialogues: list[str],
        for_generation: bool,
        detect_all: bool,
        padding: dict[str, int],
        tokenizer_type: ModelType,
    ) -> Bool[Tensor, "num_dialogues seq"] | None:
        if for_generation:
            return None
        elif detect_all:
            return torch.ones_like(tokens).bool()

        detection_mask = torch.zeros_like(tokens).bool()
        for batch_idx, dialogue in enumerate(dialogues):
            char_idx: int = 0

            for message in dialogue:
                # find the start and end of the message in the formatted text
                match = re.match(TokenizedDataset._prefix_pattern(tokenizer_type), formatted_dialogues[batch_idx][char_idx:])
                assert match is not None, f"text={formatted_dialogues[batch_idx][char_idx:]}"
                start_char_idx: int = char_idx + match.end()
                end_char_idx: int = start_char_idx + len(message.content)

                if formatted_dialogues[batch_idx][start_char_idx:end_char_idx] != message.content:
                    print(f"DEBUG: Message mismatch for {message.role}")
                    print(f"  formatted_dialogue: {repr(formatted_dialogues[batch_idx])}")
                    print(f"  char_idx: {char_idx}")
                    print(f"  match.group(): {repr(match.group())}")
                    print(f"  match.end(): {match.end()}")
                    print(f"  start_char_idx: {start_char_idx}")
                    print(f"  end_char_idx: {end_char_idx}")
                    print(f"  expected content: {repr(message.content)}")
                    print(f"  extracted content: {repr(formatted_dialogues[batch_idx][start_char_idx:end_char_idx])}")
                    
                assert (formatted_dialogues[batch_idx][start_char_idx:end_char_idx] == message.content), (formatted_dialogues[batch_idx][start_char_idx:end_char_idx], message.content)

                # use tokenizer utilities to find the corresponding tokens
                if message.detect:
                    start_tok_index = tokenizer_out.char_to_token(batch_idx, start_char_idx)
                    assert (
                        start_tok_index is not None
                    ), f"start_char_idx={start_char_idx} text={message.content}, {dialogue}"
                    
                    # Get the initial end token index using standard logic
                    end_tok_index = tokenizer_out.char_to_token(batch_idx, end_char_idx)
                    
                    # Special handling for Qwen tokenizers: manually extend by 1 token for last assistant message
                    if tokenizer_type == "qwen":
                        # Check if this is the last assistant message in the dialogue
                        is_last_assistant = (message.role == "assistant" and message == dialogue[-1])
                        
                        if is_last_assistant and end_tok_index is not None:
                            # For Qwen's last assistant message, extend the detection mask by one more token
                            # to include the last token that's often missed due to char_to_token mapping issues
                            end_tok_index += 1
                            # print(f"DEBUG Qwen: Extended detection mask for last assistant message to token {end_tok_index}")
                    
                    
                    # we also optionally include the activations that predict the first tok of
                    # the message.
                    detection_mask[
                        batch_idx,
                        start_tok_index - padding["left"] : end_tok_index + padding["right"],
                    ] = True

                # update our current position in the string
                char_idx = end_char_idx

        return detection_mask

    @classmethod
    def from_dialogue_list(
        cls,
        dialogues: list[Dialogue],
        tokenizer: PreTrainedTokenizerBase,
        padding: dict[str, dict[str, int]],
        for_generation: bool = False,
        detect_all: bool = False,
        detect_only_start_of_turn: bool = False,
        detect_only_last_token: bool = False,
        max_length: int | None = None,
        dataset_partial_id: str | None = None,
        trim_reasoning: bool = False,
        split_on: str | None = None,
        ensure_system_prompt: bool = False,
    ) -> Self:
        """
        Create an TokenizedDataset of a list of dialogues

        Args:
            dialogues (list[Dialogue]): List of dialogues to tokenize.
            tokenizer (PreTrainedTokenizerBase): Tokenizer to use for processing the text.
            padding (dict[str, dict[str, int]]): Padding configuration for detection mask. Non-zero
                values will extend the detection mask to include tokens just before or after the
                marked ones.
            for_generation (bool): Whether to have the tokens ready for generation.
                Adds the generation prompt if no assistant prompt present, else trims off the
                assistant end of turn.
            detect_all (bool): Whether to detect on all tokens rather than just marked ones.
            detect_only_start_of_turn (bool): Whether to have the detect-mask only be true at
                <|start_of_turn|> tokens.
            detect_only_last_token (bool): Whether to only detect at the last token of each message,
                the position just before <|eot_id|> tokens.
            max_length (int | None): Maximum sequence length for truncation. None means no
                truncation.
        """
        assert dialogues, "Cannot tokenize an empty dialogue"

        # Extract tokenizer type from path (handles both original and merged model paths)
        tokenizer_path_lower = tokenizer.name_or_path.lower()
        if "mistral" in tokenizer_path_lower:
            tokenizer_type = "mistral"
        elif "llama" in tokenizer_path_lower:
            tokenizer_type = "llama"
        elif "qwen" in tokenizer_path_lower:
            tokenizer_type = "qwen"
        elif "gemma" in tokenizer_path_lower:
            tokenizer_type = "gemma"
        else:
            raise NotImplementedError(f"Only gemma, mistral, llama, and qwen tokenizers are supported {tokenizer.name_or_path}")

        # if llama and no system prompt, throw a warning
        if (
            ensure_system_prompt
            and tokenizer_type == "llama"
            and not any(d[0].role == "system" for d in dialogues)
        ):
            logger.warning(
                f"No system prompts for {dataset_partial_id}. "
                f"Adding a default. e.g >>{dialogues[0]}<<"
            )
            dialogues = [
                [Message("system", "You are an AI assistant.", False)] + d for d in dialogues
            ]

        # gemma doesn't have a system prompt, so we fold it into the user prompt  
        fold_system = tokenizer_type in ["gemma"]
        processed = [preprocess_dialogue(d, fold_system=fold_system) for d in dialogues]

        final_assistants = [cls._final_message_is_assistant(d) for d in dialogues]
        assert all(final_assistants[0] == fa for fa in final_assistants)
        final_assistant = final_assistants[0]
        if for_generation:
            formatted_dialogues: list[str] = tokenizer.apply_chat_template(processed, tokenize=False, add_generation_prompt=not final_assistant)  # type: ignore
            if final_assistant:
                eot_suffix = {
                    "gemma": "<end_of_turn>\n",
                    "mistral": "</s>",
                    "llama": "<|eot_id|>",
                    "qwen": "<|im_end|>",
                }[tokenizer_type]
                formatted_dialogues = [d.removesuffix(eot_suffix) for d in formatted_dialogues]
        else:
            formatted_dialogues: list[str] = tokenizer.apply_chat_template(processed, tokenize=False, add_generation_prompt=False)  # type: ignore

        # For the following models tokenizer.apply_chat_template adds a system prompt when there is none. We want to add them therefore also to the original dialogues
        if tokenizer_type == "qwen":
            dialogues = get_dialogues_with_updated_system_prompts(formatted_dialogues, dialogues, "<|im_start|>system\n", "<|im_end|>")
        elif tokenizer_type == "mistral":
            dialogues = get_dialogues_with_updated_system_prompts(formatted_dialogues, dialogues, "[SYSTEM_PROMPT]", "[/SYSTEM_PROMPT]")

        tokenizer_out = tokenizer(
            formatted_dialogues,
            padding=True,
            return_tensors="pt",
            add_special_tokens=False,
            max_length=max_length,
            truncation=True,
        )

        tokens: torch.Tensor = tokenizer_out["input_ids"]  # type: ignore
        detection_mask = cls._get_detection_mask(
            tokens,
            tokenizer_out,
            dialogues,
            formatted_dialogues,
            for_generation,
            detect_all,
            padding[tokenizer_type],
            tokenizer_type,
        )
        if detection_mask is not None and detect_only_start_of_turn:
            sot_tok = cast(int, tokenizer("<start_of_turn>", add_special_tokens=False).input_ids[0])
            detection_mask = detection_mask & (tokens == sot_tok)
        if detection_mask is not None and detect_only_last_token:
            eot_tok_str = {
                "gemma": "<end_of_turn>\n",
                "mistral": "</s>",
                "llama": "<|eot_id|>",
                "qwen": "<|im_end|>",
            }[tokenizer_type]
            eot_tok = cast(int, tokenizer(eot_tok_str, add_special_tokens=False).input_ids[0])
            # we only want to use the detection mask on the last token of the message
            # that is, if EOT is the next token the model generated.
            next_tok = tokens.roll(shifts=-1, dims=1)
            detection_mask = detection_mask & (next_tok == eot_tok)
            # this will fail for repe and other datasets where the detect mask doesn't go until EOT

        if detection_mask is not None and detection_mask.sum() == 0:
            logger.warning("No tokens in detection mask")

        return cls(
            dialogues=dialogues,
            formatted_dialogues=formatted_dialogues,
            tokens=tokens,
            str_tokens=get_str_tokens(tokens, tokenizer),
            detection_mask=detection_mask,
            attention_mask=cast(Tensor, tokenizer_out["attention_mask"]),
            dataset_partial_id=dataset_partial_id,
        )

    @classmethod
    def from_dataset(
        cls,
        dataset: DialogueDataset,
        tokenizer: PreTrainedTokenizerBase,
        for_generation: bool = False,
        detect_all: bool = False,
        detect_only_start_of_turn: bool = False,
        detect_only_last_token: bool = False,
        max_length: int | None = None,
        additional_id: str | None = None,
    ) -> Self:
        """
        Create a TokenizedDataset from a DialogueDataset

        Args:
            dataset (DialogueDataset): The dataset containing dialogues to tokenize.
            tokenizer (PreTrainedTokenizerBase): Tokenizer to use for processing the text.
            padding (dict[str, dict[str, int]]): Padding configuration.
            for_generation (bool): Whether to prepare tokens for generation.
            detect_all (bool): Whether to detect all tokens.
            detect_only_start_of_turn (bool): Whether to detect only the start of turn tokens.
            detect_only_last_token (bool): Whether to detect only the last token before EOT.
            max_length (int | None): Maximum length of tokenized sequences.

        Returns:
            Self: A new instance of TokenizedDataset.
        """
        dataset_partial_id = f"{dataset.base_name}__{dataset.variant}"
        if additional_id:
            dataset_partial_id += f"_{additional_id}"
        return cls.from_dialogue_list(
            dialogues=dataset.dialogues,
            tokenizer=tokenizer,
            padding=dataset.padding,
            for_generation=for_generation,
            detect_all=detect_all,
            detect_only_start_of_turn=detect_only_start_of_turn,
            detect_only_last_token=detect_only_last_token,
            max_length=max_length,
            dataset_partial_id=dataset_partial_id,
        )

    def __len__(self) -> int:
        return len(self.tokens)

    def __getitem__(self, idx: int | slice) -> Self:
        idx = idx if isinstance(idx, slice) else slice(idx, idx + 1)
        return self.__class__(
            dialogues=self.dialogues[idx],
            formatted_dialogues=self.formatted_dialogues[idx],
            tokens=self.tokens[idx],
            str_tokens=self.str_tokens[idx],
            detection_mask=self.detection_mask[idx] if self.detection_mask is not None else None,
            attention_mask=self.attention_mask[idx],
        )

    def display_detection_mask(self, idx: int, tokenizer: PreTrainedTokenizerBase) -> None:
        """Display the detection mask for a given dialogue."""
        assert self.detection_mask is not None, "Detection mask not available"
        readable_tokens: list[str] = tokenizer.batch_decode(
            self.tokens[idx], clean_up_tokenization_spaces=False
        )
        for readable_token, mask in zip(readable_tokens, self.detection_mask[idx], strict=True):
            readable_token = readable_token.replace("\n", "\\n")
            color = "red" if mask else "light_grey"
            print(colored(readable_token, color), end="")
            print("|", end="")


def get_tokenized_dataset_pair(
    honest: DialogueDataset,
    deceptive: DialogueDataset,
    tokenizer: PreTrainedTokenizerBase,
    enforce_equal_detect_lengths: bool = False,
    detect_all: bool = False,
    detect_only_start_of_turn: bool = False,
    detect_only_last_token: bool = False,
) -> tuple[TokenizedDataset, TokenizedDataset]:
    """
    Tokenizes the input dialogues in a PairedDialogueDataset, returning a pair of TokenizedDatasets.

    Args:
        honest (DialogueDataset): Dataset containing truthful dialogues.
        deceptive (DialogueDataset): Dataset containing deceptive dialogues.
        tokenizer (PreTrainedTokenizerBase): Tokenizer to use for processing the text.
        enforce_equal_detect_lengths (bool): If True, this will trim the detection mask to the
            shorter of the pair.
        detect_all (bool): Whether to detect all tokens.
        detect_only_start_of_turn (bool): Whether to detect only the start of turn tokens.
        detect_only_last_token (bool): Whether to detect only the last token before EOT.

    Returns:
        tuple[TokenizedDataset, TokenizedDataset]: A pair of TokenizedDataset instances.
    """

    truthful_tokenized = TokenizedDataset.from_dataset(
        honest,
        tokenizer,
        detect_all=detect_all,
        detect_only_start_of_turn=detect_only_start_of_turn,
        detect_only_last_token=detect_only_last_token,
        additional_id="honest",
    )
    deceptive_tokenized = TokenizedDataset.from_dataset(
        deceptive,
        tokenizer,
        detect_all=detect_all,
        detect_only_start_of_turn=detect_only_start_of_turn,
        detect_only_last_token=detect_only_last_token,
        additional_id="deceptive",
    )

    if enforce_equal_detect_lengths:
        equalize_detection_masks(truthful_tokenized, deceptive_tokenized)

    return truthful_tokenized, deceptive_tokenized


def equalize_detection_masks(dataset1: TokenizedDataset, dataset2: TokenizedDataset):
    """
    Modifies the detection masks of both datasets in place to ensure they have equal lengths.
    """
    assert dataset1.detection_mask is not None
    assert dataset2.detection_mask is not None
    for i in range(len(dataset1.detection_mask)):
        mask1 = dataset1.detection_mask[i]
        mask2 = dataset2.detection_mask[i]

        start1, end1 = get_mask_range(mask1)
        start2, end2 = get_mask_range(mask2)

        length1 = end1 - start1
        length2 = end2 - start2
        min_length = min(length1, length2)

        # Trim the detection masks to the minimum length
        dataset1.detection_mask[i][start1 + min_length : end1] = False
        dataset2.detection_mask[i][start2 + min_length : end2] = False


def get_mask_range(mask: torch.Tensor) -> tuple[int, int]:
    """
    Finds the start and end indices of the True values in the mask.
    """
    start = torch.argmax((mask == True).int())
    end = start + torch.argmax((mask[start:] == False).int())

    # If no end found, set end to the length of the mask
    if end == start:
        end = len(mask)

    start = int(start)
    end = int(end)

    return start, end
