"""
Dataset module -- Member A (Huy)
PyTorch Dataset for prompt classification (Jailbreak vs Benign).
"""
import torch
from torch.utils.data import Dataset


class PromptDataset(Dataset):
    """
    Dataset cho bai toan phan loai prompt jailbreak.
    - texts: list of raw text strings
    - labels: list of int (0=benign, 1=jailbreak) or str labels; None for inference
    - tokenizer: HuggingFace tokenizer; None returns raw text dict
    - max_length: max token length for truncation
    """
    LABEL_MAP = {
        "benign": 0,
        "jailbreak": 1,
        "prompt_injection": 2,
    }
    ID2LABEL = {0: "BENIGN", 1: "JAILBREAK", 2: "PROMPT_INJECTION"}
    LABEL2ID = {"BENIGN": 0, "JAILBREAK": 1, "PROMPT_INJECTION": 2}

    def __init__(self, texts, labels=None, tokenizer=None, max_length: int = 256):
        self.texts = [str(t) for t in texts]
        self.tokenizer = tokenizer
        self.max_length = max_length

        if labels is not None:
            processed_labels = []
            for lab in labels:
                if isinstance(lab, str):
                    processed_labels.append(self.LABEL_MAP.get(lab.lower(), 0))
                else:
                    processed_labels.append(int(lab))
            self.labels = processed_labels
        else:
            self.labels = None

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

    def __getitem__(self, idx: int) -> dict:
        text = self.texts[idx]

        if self.tokenizer is not None:
            encoding = self.tokenizer(
                text,
                truncation=True,
                max_length=self.max_length,
                padding="max_length",
                return_tensors="pt",
            )
            item = {
                "input_ids": encoding["input_ids"].squeeze(0),
                "attention_mask": encoding["attention_mask"].squeeze(0),
            }
            if "token_type_ids" in encoding:
                item["token_type_ids"] = encoding["token_type_ids"].squeeze(0)
        else:
            item = {"text": text}

        if self.labels is not None:
            item["label"] = torch.tensor(self.labels[idx], dtype=torch.long)

        return item