"""Local JSONL datasets for multimodal VLM training."""

from typing import Dict
import torch
from torch.utils.data import Dataset

from taoTrain.data.vlm_utils import (
    build_image_transform,
    build_vlm_sequence_tokens,
    load_image,
    load_multimodal_records,
    load_tokenizer_from_config,
    parse_vlm_record,
    validate_vlm_special_tokens,
)


class VLMJSONLDataset(Dataset):
    """JSONL dataset for multimodal connector training and multimodal SFT."""

    def __init__(self, config, split: str = "train"):
        """Initialize the multimodal dataset."""
        self.config = config
        self.split = split
        self.records = load_multimodal_records(config)
        self.tokenizer = load_tokenizer_from_config(config)
        self.special_token_ids = validate_vlm_special_tokens(config, self.tokenizer)
        self.transform = build_image_transform(config.image_size)
        self.text_seq_length = config.model.max_seq_length - config.vision_prefix_tokens + 1

        if self.text_seq_length < 2:
            raise ValueError(
                "model.max_seq_length must be at least vision_prefix_tokens + 1 for multimodal expansion"
            )

    def __len__(self) -> int:
        """Return dataset size."""
        return len(self.records)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """Load one multimodal training example."""
        record = self.records[idx]
        image_path, turns = parse_vlm_record(record, self.config)
        pixel_values = load_image(image_path, self.config.dataset.jsonl_path, self.transform)

        input_ids, attention_mask, mask, image_token_id = build_vlm_sequence_tokens(
            turns=turns,
            tokenizer=self.tokenizer,
            image_token=self.config.image_token,
            user_token=self.config.user_token,
            assistant_token=self.config.assistant_token,
            max_seq_length=self.text_seq_length,
        )

        first_non_pad_idx = next((i for i, value in enumerate(attention_mask) if value == 1), None)
        if first_non_pad_idx != 0 or input_ids[0] != image_token_id:
            raise ValueError("Multimodal samples must begin with the configured <image> token")

        labels = input_ids[1:].copy() + [-100]
        for token_idx, mask_value in enumerate(mask):
            if mask_value == 0:
                labels[token_idx] = -100

        if all(label == -100 for label in labels):
            raise ValueError(
                f"VLM sample at index {idx} produced no trainable assistant/caption tokens. "
                "Check the record format and special-token masking."
            )

        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
            "pixel_values": pixel_values,
        }


class VLMSFTJSONLDataset(VLMJSONLDataset):
    """JSONL dataset for end-to-end multimodal supervised fine-tuning."""

    pass