"""BERT fine-tuning and embedding extraction utilities."""

from __future__ import annotations

from typing import Any


def BERTDataset(
    texts: list[str],
    labels: list[int] | None = None,
    tokenizer: Any = None,
    max_length: int = 128,
    model_name: str = "distilbert-base-uncased",
) -> Any:
    """PyTorch Dataset for BERT-style input (input_ids, attention_mask, optional labels). Requires transformers and torch."""
    try:
        import torch
        from torch.utils.data import Dataset
        from transformers import AutoTokenizer
    except ImportError as e:
        raise ImportError(
            "BERTDataset requires torch and transformers. Install with: pip install torch transformers"
        ) from e
    if tokenizer is None:
        tokenizer = AutoTokenizer.from_pretrained(model_name)

    class _BERTDataset(Dataset):
        def __init__(self):
            self.encodings = tokenizer(
                texts,
                truncation=True,
                padding="max_length",
                max_length=max_length,
                return_tensors="pt",
            )
            self.labels = torch.tensor(labels, dtype=torch.long) if labels is not None else None

        def __len__(self) -> int:
            return self.encodings["input_ids"].size(0)

        def __getitem__(self, i: int) -> Any:
            item = {
                "input_ids": self.encodings["input_ids"][i],
                "attention_mask": self.encodings["attention_mask"][i],
            }
            if self.labels is not None:
                item["labels"] = self.labels[i]
            return item

    return _BERTDataset()


def get_bert_tokenizer(model_name: str = "distilbert-base-uncased") -> Any:
    """Return HuggingFace tokenizer for the given model. Requires transformers."""
    try:
        from transformers import AutoTokenizer
        return AutoTokenizer.from_pretrained(model_name)
    except ImportError as e:
        raise ImportError(
            "get_bert_tokenizer requires transformers. Install with: pip install transformers"
        ) from e
