"""
DataLoader implementation for ASR training.
""" 

import torch
from torch.utils.data import DataLoader
from transformers import WhisperProcessor

class ASRDataLoader:
    def __init__(
        self,
        processor: WhisperProcessor,
        batch_size: int = 32,
        num_workers: int = 2,
        max_frames: int = 3000,
        sample_rate: int = 16000,
        pin_memory: bool = True,
        persistent_workers: bool = True,
        prefetch_factor: int = 2
    ):

        self.processor = processor
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.max_frames = max_frames
        self.sample_rate = sample_rate
        self.pin_memory = pin_memory
        self.persistent_workers = persistent_workers
        self.prefetch_factor = prefetch_factor

    def collate_fn(self, batch):
        batch = [b for b in batch if b and b["transcription"].strip()]
        if not batch:
            return None

        processed_features = []
        processed_labels = []

        for item in batch:
            audio_array = item["audio"]["array"]
            text = item["transcription"]

            audio_inputs = self.processor.feature_extractor(
                audio_array,
                sampling_rate=self.sample_rate,
                return_tensors="pt"
            )

            text_inputs = self.processor.tokenizer(
                text,
                padding="max_length",
                truncation=True,
                max_length=256,
                return_tensors="pt"
            )

            feats = audio_inputs.input_features[0]
            if feats.shape[-1] > self.max_frames:
                feats = feats[..., :self.max_frames]
            elif feats.shape[-1] < self.max_frames:
                pad_size = self.max_frames - feats.shape[-1]
                feats = torch.nn.functional.pad(feats, (0, pad_size))

            processed_features.append(feats)

            labels = text_inputs.input_ids[0]
            labels[labels == self.processor.tokenizer.pad_token_id] = -100
            processed_labels.append(labels)

        batch_features = torch.stack(processed_features, dim=0)
        batch_labels = torch.stack(processed_labels, dim=0)

        return {
            "input_features": batch_features,
            "labels": batch_labels
        }
        
    def safe_collate_fn(self, batch):
        try:
            result = self.collate_fn(batch)
            if result is None:
                return {
                    "input_features": torch.zeros(1, 80, self.max_frames),
                    "labels": torch.full((1, 256), -100, dtype=torch.long)
                }
            return result
        except Exception as e:
            print(f"Collate error: {e}")
            return {
                "input_features": torch.zeros(1, 80, self.max_frames),
                "labels": torch.full((1, 256), -100, dtype=torch.long)
            }
        
    def get_loader(self, dataset, shuffle=True, drop_last=True):
        return DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=shuffle,
            num_workers=self.num_workers,
            pin_memory=self.pin_memory,
            collate_fn=self.safe_collate_fn,
            persistent_workers=self.persistent_workers,
            prefetch_factor=self.prefetch_factor,
            drop_last=drop_last
        )