"""PyTorch Dataset & DataLoader for ASVspoof 5 with on-the-fly telephony degradation."""

from typing import Dict, Optional, Tuple, Union
import os
import torch
from torch.utils.data import Dataset, DataLoader
import soundfile as sf
import pandas as pd

from src.transforms import TelephonyDegradation, TandemTelephonyDegradation


def pad_or_crop_waveform(
    waveform: torch.Tensor,
    target_len: int,
    is_train: bool = True
) -> torch.Tensor:
    """Pad or crop waveform tensor (1, N) to target_len samples."""
    num_samples = waveform.shape[-1]
    if num_samples == target_len:
        return waveform
    elif num_samples > target_len:
        if is_train:
            max_start = num_samples - target_len
            start = torch.randint(0, max_start + 1, (1,)).item()
        else:
            start = (num_samples - target_len) // 2
        return waveform[..., start:start + target_len]
    else:
        # Repeat or zero pad
        repeat_factor = (target_len // num_samples) + 1
        repeated = waveform.repeat(1, repeat_factor)
        return repeated[..., :target_len]


class TelephonyASVDataset(Dataset):
    """ASVspoof 5 Dataset with on-the-fly telephony degradation.
    
    Reads manifest CSV, loads 16 kHz raw audio, pads/crops to fixed duration,
    and applies on-the-fly 8 kHz telephony degradation (supports single or dual tandem compression).
    """

    def __init__(
        self,
        manifest_path: str,
        duration_sec: float = 4.0,
        is_train: bool = True,
        apply_telephony: bool = True,
        use_tandem: bool = True,
        voip_codec: str = "ULAW",
        packet_loss_rate: float = 0.05,
    ):
        super().__init__()
        if not os.path.exists(manifest_path):
            raise FileNotFoundError(f"Manifest file not found: {manifest_path}")

        self.df = pd.read_csv(manifest_path)
        self.duration_sec = duration_sec
        self.is_train = is_train
        self.apply_telephony = apply_telephony
        self.use_tandem = use_tandem

        # 16 kHz target length in raw audio
        self.target_raw_samples = int(16000 * duration_sec)

        # On-the-fly degradation module: dual tandem (Cellular GSM -> VoIP ULAW/Opus) or standard
        if use_tandem:
            self.degradation = TandemTelephonyDegradation(
                orig_sample_rate=16000,
                target_sample_rate=8000,
                cellular_codec="GSM610",
                voip_codec=voip_codec,
                packet_loss_rate=packet_loss_rate if is_train else 0.0,
            )
        else:
            self.degradation = TelephonyDegradation(
                orig_sample_rate=16000,
                target_sample_rate=8000,
                apply_bandpass=True,
                apply_mulaw=True,
                packet_loss_rate=packet_loss_rate if is_train else 0.0,
            )

        if is_train:
            self.degradation.train()
        else:
            self.degradation.eval()

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

    def __getitem__(self, idx: int) -> Dict[str, Union[torch.Tensor, str, int]]:
        row = self.df.iloc[idx]
        file_path = row["file_path"]
        label = int(row["label"])  # 0: bonafide, 1: spoof
        speaker_id = str(row["speaker_id"])
        file_id = str(row["file_id"])

        # 1. Load audio with soundfile (fast & robust)
        audio_data, sr = sf.read(file_path)
        if sr != 16000:
            raise ValueError(f"Expected 16000 Hz audio, found {sr} Hz in {file_path}")

        # Ensure tensor shape (1, num_samples)
        waveform = torch.from_numpy(audio_data).float()
        if waveform.ndim == 1:
            waveform = waveform.unsqueeze(0)
        elif waveform.ndim == 2 and waveform.shape[0] > 1:
            # Multi-channel -> mono
            waveform = waveform.mean(dim=0, keepdim=True)

        # 2. Pad or crop to target raw duration (e.g. 4.0s @ 16kHz = 64,000 samples)
        waveform = pad_or_crop_waveform(waveform, self.target_raw_samples, is_train=self.is_train)

        # 3. Apply on-the-fly telephony degradation (returns 8kHz = 32,000 samples)
        if self.apply_telephony:
            with torch.no_grad():
                processed_audio = self.degradation(waveform)
        else:
            processed_audio = waveform

        return {
            "audio": processed_audio.squeeze(0),  # (num_samples,) e.g. 32000 @ 8kHz
            "label": torch.tensor(label, dtype=torch.float32),
            "speaker_id": speaker_id,
            "file_id": file_id,
        }


def get_dataloader(
    manifest_path: str,
    batch_size: int = 16,
    is_train: bool = True,
    duration_sec: float = 4.0,
    apply_telephony: bool = True,
    use_tandem: bool = True,
    voip_codec: str = "ULAW",
    packet_loss_rate: float = 0.05,
    num_workers: int = 0,
    pin_memory: bool = False,
) -> DataLoader:
    """Convenience factory for TelephonyASVDataset DataLoader."""
    dataset = TelephonyASVDataset(
        manifest_path=manifest_path,
        duration_sec=duration_sec,
        is_train=is_train,
        apply_telephony=apply_telephony,
        use_tandem=use_tandem,
        voip_codec=voip_codec,
        packet_loss_rate=packet_loss_rate,
    )

    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=is_train,
        num_workers=num_workers,
        pin_memory=pin_memory,
        drop_last=is_train,
    )
    return loader