"""Data collators for instruction tuning datasets."""

import torch
from dataclasses import dataclass
from typing import Dict, List, Any, Optional
from composennent.nlp.tokenizers.base import BaseTokenizer


@dataclass
class InstructionCollator:
    """Basic collator for instruction tuning datasets.

    This collator:
    1. Tokenizes text samples
    2. Pads them to the same length
    3. Creates labels for training
    4. Creates attention masks

    Args:
        tokenizer: Tokenizer instance
        max_length: Maximum sequence length (default: 2048)
        padding: Padding strategy - "longest" or "max_length" (default: "longest")
        ignore_index: Index to use for masked tokens in labels (default: -100)

    Example:
        >>> from torch.utils.data import DataLoader
        >>> collator = InstructionCollator(tokenizer, max_length=2048)
        >>> dataloader = DataLoader(dataset, batch_size=4, collate_fn=collator)
    """

    tokenizer: BaseTokenizer
    max_length: int = 2048
    padding: str = "longest"
    ignore_index: int = -100

    def __call__(self, batch: List[Dict[str, str]]) -> Dict[str, torch.Tensor]:
        """Collate a batch of examples into tensors.

        Args:
            batch: List of dictionaries with 'text' key

        Returns:
            Dictionary with:
                - input_ids: Tokenized and padded input (batch_size, seq_len)
                - attention_mask: Mask for valid tokens (batch_size, seq_len)
                - labels: Labels for training (batch_size, seq_len)
        """
        texts = [item["text"] for item in batch]


        if self.padding == "max_length":
            pad_length = self.max_length
        else:

            all_tokens = [self.tokenizer.encode(text) for text in texts]
            pad_length = min(max(len(t) for t in all_tokens), self.max_length)


        input_ids_list = []
        attention_mask_list = []

        for text in texts:
            tokens = self.tokenizer.encode(text)


            if len(tokens) > self.max_length:
                tokens = tokens[:self.max_length]


            attention_mask = [1] * len(tokens)


            padding_length = pad_length - len(tokens)
            if padding_length > 0:
                tokens = tokens + [self.tokenizer.pad_token_id] * padding_length
                attention_mask = attention_mask + [0] * padding_length

            input_ids_list.append(tokens)
            attention_mask_list.append(attention_mask)


        input_ids = torch.tensor(input_ids_list, dtype=torch.long)
        attention_mask = torch.tensor(attention_mask_list, dtype=torch.long)


        labels = input_ids.clone()


        labels[labels == self.tokenizer.pad_token_id] = self.ignore_index

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": labels,
        }


@dataclass
class InstructionCollatorWithPromptMasking(InstructionCollator):
    """Advanced collator that masks the prompt in labels.

    This collator trains the model only on the response portion,
    not on the instruction/input. This is the recommended approach
    for instruction tuning as it prevents the model from learning
    to generate instructions.

    Supports 5 popular instruction formats:
        - "chatml": OpenAI ChatML format - `<|im_start|>assistant\\n`
        - "alpaca": Stanford Alpaca format - `### Response:\\n`
        - "llama": Llama 2/3 Chat format - `[/INST]`
        - "vicuna": Vicuna/FastChat format - `ASSISTANT:`
        - "zephyr": HuggingFace Zephyr format - `<|assistant|>\\n`

    Args:
        tokenizer: Tokenizer instance
        max_length: Maximum sequence length (default: 2048)
        padding: Padding strategy - "longest" or "max_length" (default: "longest")
        ignore_index: Index to use for masked tokens in labels (default: -100)
        format_type: One of "chatml", "alpaca", "llama", "vicuna", "zephyr" (default: "chatml")
        response_template: Custom response template. If provided, overrides format_type.

    Example:
        >>> # Using preset format
        >>> collator = InstructionCollatorWithPromptMasking(
        ...     tokenizer,
        ...     format_type="alpaca"
        ... )
        >>> 
        >>> # Using custom template
        >>> collator = InstructionCollatorWithPromptMasking(
        ...     tokenizer,
        ...     response_template="### Answer:\\n"
        ... )
        >>> dataloader = DataLoader(dataset, collate_fn=collator)
    """

    # Preset response templates for popular formats
    RESPONSE_TEMPLATES = {
        "chatml": "<|im_start|>assistant\n",
        "alpaca": "### Response:\n",
        "llama": "[/INST]",
        "vicuna": "ASSISTANT:",
        "zephyr": "<|assistant|>\n",
    }

    format_type: str = "chatml"
    response_template: Optional[str] = None

    def __post_init__(self):
        """Set response_template based on format_type if not provided."""
        if self.response_template is None:
            if self.format_type not in self.RESPONSE_TEMPLATES:
                valid_formats = ", ".join(f'"{k}"' for k in self.RESPONSE_TEMPLATES.keys())
                raise ValueError(
                    f"Unknown format_type: '{self.format_type}'. "
                    f"Choose from: {valid_formats}, or provide a custom response_template."
                )
            self.response_template = self.RESPONSE_TEMPLATES[self.format_type]

    def __call__(self, batch: List[Dict[str, str]]) -> Dict[str, torch.Tensor]:
        """Collate batch with prompt masking.

        Args:
            batch: List of dictionaries with 'text' key

        Returns:
            Dictionary with input_ids, attention_mask, and labels (prompt masked)
        """
        texts = [item["text"] for item in batch]


        if self.padding == "max_length":
            pad_length = self.max_length
        else:
            all_tokens = [self.tokenizer.encode(text) for text in texts]
            pad_length = min(max(len(t) for t in all_tokens), self.max_length)


        input_ids_list = []
        attention_mask_list = []
        labels_list = []

        for text in texts:
            tokens = self.tokenizer.encode(text)


            if len(tokens) > self.max_length:
                tokens = tokens[:self.max_length]


            labels = tokens.copy()


            response_start_idx = text.find(self.response_template)

            if response_start_idx != -1:

                prompt_text = text[:response_start_idx + len(self.response_template)]
                prompt_tokens = self.tokenizer.encode(prompt_text)
                prompt_length = min(len(prompt_tokens), len(tokens))


                labels[:prompt_length] = [self.ignore_index] * prompt_length
            else:


                labels = [self.ignore_index] * len(tokens)


            attention_mask = [1] * len(tokens)


            padding_length = pad_length - len(tokens)
            if padding_length > 0:
                tokens = tokens + [self.tokenizer.pad_token_id] * padding_length
                attention_mask = attention_mask + [0] * padding_length
                labels = labels + [self.ignore_index] * padding_length

            input_ids_list.append(tokens)
            attention_mask_list.append(attention_mask)
            labels_list.append(labels)


        input_ids = torch.tensor(input_ids_list, dtype=torch.long)
        attention_mask = torch.tensor(attention_mask_list, dtype=torch.long)
        labels = torch.tensor(labels_list, dtype=torch.long)

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": labels,
        }



@dataclass
class InferenceCollator:
    """Collator for inference (no labels needed).

    Use this when generating text, not training.

    Args:
        tokenizer: Tokenizer instance
        max_length: Maximum sequence length (default: 2048)
        padding: Padding strategy (default: "longest")

    Example:
        >>> collator = InferenceCollator(tokenizer)
        >>> dataloader = DataLoader(dataset, collate_fn=collator)
        >>> for batch in dataloader:
        ...     outputs = model.generate(**batch)
    """

    tokenizer: BaseTokenizer
    max_length: int = 2048
    padding: str = "longest"

    def __call__(self, batch: List[Dict[str, str]]) -> Dict[str, torch.Tensor]:
        """Collate batch for inference.

        Args:
            batch: List of dictionaries with 'text' key

        Returns:
            Dictionary with input_ids and attention_mask only
        """
        texts = [item["text"] for item in batch]


        if self.padding == "max_length":
            pad_length = self.max_length
        else:
            all_tokens = [self.tokenizer.encode(text) for text in texts]
            pad_length = min(max(len(t) for t in all_tokens), self.max_length)


        input_ids_list = []
        attention_mask_list = []

        for text in texts:
            tokens = self.tokenizer.encode(text)


            if len(tokens) > self.max_length:
                tokens = tokens[:self.max_length]


            attention_mask = [1] * len(tokens)


            padding_length = pad_length - len(tokens)
            if padding_length > 0:
                tokens = tokens + [self.tokenizer.pad_token_id] * padding_length
                attention_mask = attention_mask + [0] * padding_length

            input_ids_list.append(tokens)
            attention_mask_list.append(attention_mask)


        input_ids = torch.tensor(input_ids_list, dtype=torch.long)
        attention_mask = torch.tensor(attention_mask_list, dtype=torch.long)

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
        }
