"""
LibriSpeech dataset wrapper + collate function.

torchaudio.datasets.LIBRISPEECH reads <root>/LibriSpeech/<split>/... and
returns waveforms shaped (1, L).
"""

from typing import cast

import torch
from torch.utils.data import ConcatDataset, Dataset
import torchaudio

from src.tokenizer import ASRTokenizer, PAD_ID

# LIBRISPEECH.__getitem__'s declared return type; ConcatDataset erases it.
LibriSpeechItem = tuple[torch.Tensor, int, str, int, int, int]


class LibriSpeechASR(Dataset):
    def __init__(self, root: str, splits, download: bool = False):
        if isinstance(splits, str):
            splits = [splits]
        self.datasets = [
            torchaudio.datasets.LIBRISPEECH(root=root, url=split, download=download)
            for split in splits
        ]
        self._concat = ConcatDataset(self.datasets)

    def __len__(self):
        return len(self._concat)

    def __getitem__(self, idx):
        waveform, sample_rate, transcript, _speaker_id, _chapter_id, _utterance_id = cast(
            LibriSpeechItem, self._concat[idx]
        )
        assert sample_rate == 16000, f"expected 16kHz audio, got {sample_rate}"
        return waveform.squeeze(0), transcript


class ASRCollate:
    """Pads a batch and produces the tensors forward_train/forward_eval need:
    waveforms, CTC targets, and decoder teacher-forcing input/target pairs."""

    def __init__(self, tokenizer: ASRTokenizer):
        self.tokenizer = tokenizer

    def __call__(self, batch):
        waveforms, transcripts = zip(*batch)

        wave_lengths = torch.tensor([w.size(0) for w in waveforms], dtype=torch.long)
        max_wave_len = int(wave_lengths.max().item())
        padded_waveforms = torch.zeros(len(waveforms), max_wave_len, dtype=torch.float32)
        for i, w in enumerate(waveforms):
            padded_waveforms[i, : w.size(0)] = w

        ctc_ids_list = [self.tokenizer.encode(t) for t in transcripts]
        ctc_target_lengths = torch.tensor([len(ids) for ids in ctc_ids_list], dtype=torch.long)
        max_ctc_len = max(1, int(ctc_target_lengths.max().item()))
        ctc_targets = torch.full((len(batch), max_ctc_len), PAD_ID, dtype=torch.long)
        for i, ids in enumerate(ctc_ids_list):
            if ids:
                ctc_targets[i, : len(ids)] = torch.tensor(ids, dtype=torch.long)

        dec_pairs = [self.tokenizer.decoder_input_and_target(t) for t in transcripts]
        decoder_input_lengths = torch.tensor([len(p[0]) for p in dec_pairs], dtype=torch.long)
        max_dec_len = int(decoder_input_lengths.max().item())
        decoder_input = torch.full((len(batch), max_dec_len), PAD_ID, dtype=torch.long)
        decoder_target = torch.full((len(batch), max_dec_len), PAD_ID, dtype=torch.long)
        for i, (inp, tgt) in enumerate(dec_pairs):
            decoder_input[i, : len(inp)] = torch.tensor(inp, dtype=torch.long)
            decoder_target[i, : len(tgt)] = torch.tensor(tgt, dtype=torch.long)

        return {
            "waveforms": padded_waveforms,
            "wave_lengths": wave_lengths,
            "ctc_targets": ctc_targets,
            "ctc_target_lengths": ctc_target_lengths,
            "decoder_input": decoder_input,
            "decoder_input_lengths": decoder_input_lengths,
            "decoder_target": decoder_target,
            "transcripts": transcripts,
        }