import random

import torch
from torch import Tensor
from torch.utils.data import Dataset, DataLoader
import torchaudio
from torch_audiomentations import ApplyImpulseResponse, Gain, PitchShift, LowPassFilter, HighPassFilter
from torchaudio.transforms import Resample, TimeStretch
from pathlib import Path
import pandas as pd
from typing import List, TypedDict
from handle.text_normalizer import collapse_spaces, normalize_extended_uyghur_characters
from tokenizer import ASRTokenizer


# 单个样本的数据结构（Dataset.__getitem__ 返回）
class BatchItem(TypedDict):
    waveform: Tensor          # [time]
    target_ids: Tensor        # [seq_len] 目标文本的token IDs
    target_text: str          # 原始文本
    audio_path: str           # 音频文件路径


# 批量数据的数据结构（collate_fn 返回，DataLoader 输出）
class Batch(TypedDict):
    waveforms: Tensor         # [batch, time]
    targets: Tensor           # [batch, max_len] padding后的目标IDs
    waveform_lengths: Tensor       # [batch] 每个样本的实际Waveform长度
    target_lengths: Tensor    # [batch] 每个样本的实际目标长度
    target_texts: List[str]   # [batch] 原始文本列表
    audio_paths: List[str]    # [batch] 音频路径列表

class TsvFormat(TypedDict):
    client_id: str
    path: str
    sentence: str
    up_votes: int
    down_votes: int
    age: str
    gender: str
    locale: str

class NoiseAugmentor:
    def __init__(self, noise_root: Path, sample_rate: int=16000):
        self.sample_rate = sample_rate
        self.noise_files = list(Path(noise_root).rglob("*.wav"))

    def apply_real_noise(self, waveform: Tensor):
        # 1. 随机选一个噪音文件
        noise_path = random.choice(self.noise_files)
        noise_waveform, sr = torchaudio.load_with_torchcodec(noise_path)
        
        # Resample to target sample rate.
        if sr != self.sample_rate:
            noise_waveform = Resample(sr, self.sample_rate)(noise_waveform)
        
        # Convert to mono if it is setro.
        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)
        
        # 3. 截取或填充，使其长度与语音一致
        sig_len = waveform.shape[1]
        noise_len = noise_waveform.shape[1]
        
        if noise_len >= sig_len:
            # 随机截取一段
            start = random.randint(0, noise_len - sig_len)
            noise_waveform = noise_waveform[:, start:start + sig_len]
        else:
            full_noise = torch.zeros_like(waveform)
            start = random.randint(0, sig_len - noise_len)
            full_noise[:, start : start + noise_len] = noise_waveform
            noise_waveform = full_noise
            
        # 4. 设定随机信噪比 SNR (5dB 到 20dB)
        snr_db = random.uniform(5, 20)
        
        # 5. 混合
        return self._mix_at_snr(waveform, noise_waveform, snr_db)

    def _mix_at_snr(self, signal: Tensor, noise: Tensor, snr_db: float):
        s_p = signal.pow(2).mean()
        n_p = noise.pow(2).mean()
        snr_linear = 10**(snr_db / 10)
        scale = torch.sqrt(s_p / (n_p * snr_linear + 1e-8))
        
        noisy = signal + scale * noise
        # 归一化，防止溢出
        return noisy / (noisy.abs().max() + 1e-8)

class CommonVoiceDataset(Dataset[BatchItem]):
    def __init__(
        self,
        tsv_path: Path,
        audio_dir: Path,
        noise_dir: Path,
        tokenizer: ASRTokenizer,
        sample_rate: int = 16000,
        max_audio_len: int = 480000,  # 30秒 @ 16kHz
        augment: bool = True,
        augment_prob: float = 0.5,
    ) -> None:
        super().__init__()
        self.noise_augmentor = NoiseAugmentor(noise_root=noise_dir, sample_rate=sample_rate)
        self.audio_dir = Path(audio_dir)
        self.tokenizer = tokenizer
        self.sample_rate = sample_rate
        self.max_audio_len = max_audio_len
        self.augment = augment
        self.augment_prob = augment_prob

        self.data: pd.DataFrame = pd.read_csv(tsv_path, sep='\t')

        valid_indices = []
        for index, row in self.data.iterrows():
            audio_path: Path = self.audio_dir / row['path']
            if audio_path.exists():
                valid_indices.append(index)

        self.data = self.data.loc[valid_indices].reset_index(drop=True)

        self.gain_up = Gain(min_gain_in_db=4, max_gain_in_db=8, p=1.0, output_type='tensor')
        self.gain_down = Gain(min_gain_in_db=-15, max_gain_in_db=-8, p=1.0, output_type='tensor')
        self.pitch_up = PitchShift(min_transpose_semitones=1, max_transpose_semitones=4, p=1.0, sample_rate=self.sample_rate, output_type='tensor')
        self.pitch_down = PitchShift(min_transpose_semitones=-4, max_transpose_semitones=-1, p=1.0, sample_rate=self.sample_rate, output_type='tensor')
        self.lowpass = LowPassFilter(min_cutoff_freq=600, max_cutoff_freq=2000, p=1.0, output_type='tensor')
        self.highpass = HighPassFilter(min_cutoff_freq=800, max_cutoff_freq=2000, p=1.0, output_type='tensor')
        self.apply_ir = ApplyImpulseResponse(ir_paths=noise_dir,convolve_mode='same', p=1, output_type="tensor")

    def __len__(self):
        return len(self.data)
    
    def _load_audio(self, audio_path: Path) -> Tensor:
        waveform, sample_rate = torchaudio.load_with_torchcodec(audio_path)

        # Resample to target sample rate.
        if sample_rate != self.sample_rate:
            waveform = Resample(sample_rate, self.sample_rate)(waveform)

        # Convert to mono if it is setro.
        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)
        
        # Normalization
        waveform = waveform / waveform.abs().max()
        
        # Clip waveform exceeds from max length.
        if waveform.shape[1] > self.max_audio_len:
            waveform = waveform[:, :self.max_audio_len]
        return waveform
    
    def _augment_waveform(self, waveform: Tensor) -> Tensor:
        if not self.augment or random.random() > self.augment_prob:
            return waveform
        
        # 1. voice Stretch/Compress 
        if random.random() < 0.6:
            waveform = self._stretch_or_compress(waveform=waveform)
        
        if random.random() < 0.4:
            waveform = self.noise_augmentor.apply_real_noise(waveform)
        
        if random.random() < 0.3:
            waveform = self._time_mask_waveform(waveform=waveform)

        # torch_audiomentations: [1, time] -> [1, 1, time]
        if waveform.dim() == 2:
            waveform_3d = waveform.unsqueeze(0)
        # 随机选择一种频谱增强
        choice = random.random()
        if choice < 0.15:
            # 增益变化（上或下）
            if random.random() < 0.5:
                waveform_3d = self.gain_up(waveform_3d, sample_rate=self.sample_rate)
            else:
                waveform_3d = self.gain_down(waveform_3d, sample_rate=self.sample_rate)
        elif choice < 0.25:
            # 音高变化（上或下）
            if random.random() < 0.5:
                waveform_3d = self.pitch_up(waveform_3d, sample_rate=self.sample_rate)
            else:
                waveform_3d = self.pitch_down(waveform_3d, sample_rate=self.sample_rate)
        elif choice < 0.30:
            # 低通滤波（声音发闷）
            waveform_3d = self.lowpass(waveform_3d, sample_rate=self.sample_rate)
        # elif choice < 0.32:
        #     # 低通滤波（声音发闷）
        #     waveform_3d = self.apply_ir(waveform_3d, sample_rate=self.sample_rate)
        elif choice < 0.35:
            # 高通滤波（电话效果）
            waveform_3d = self.highpass(waveform_3d, sample_rate=self.sample_rate)
        
        #  [1, 1, time] -> [1, time]
        waveform = waveform_3d.squeeze(0)

        # 防止多次 augment 后振幅溢出，最后归一化
        max_amp = waveform.abs().max()
        if max_amp > 1.0:
            waveform = waveform / max_amp
        
        return waveform
        
    def _stretch_or_compress(self, waveform: Tensor) -> Tensor:
        speed_factor = random.uniform(0.85, 1.4) # (Speed Change: 0.85x - 1.4x)
        spec = torch.stft(
            waveform.squeeze(0),
            n_fft=400,
            hop_length=160,
            window=torch.hann_window(400).to(waveform.device),
            return_complex=True
        )
        
        # 时间拉伸（不改变音高）
        stretched_spec = TimeStretch(hop_length=160, n_freq=spec.shape[-2], fixed_rate=speed_factor)(spec)
        
        # 转回波形
        waveform_stretched = torch.istft(stretched_spec, n_fft=400, hop_length=160, window=torch.hann_window(400).to(waveform.device)).unsqueeze(0)
    
        return waveform_stretched
    
    def _time_mask_waveform(self, waveform: Tensor) -> Tensor:
        audio_len = waveform.shape[1]
        sr = self.sample_rate # 16000
        
        # 设置参数：单次遮盖最长 0.4 秒 (6400个点)
        max_mask_time = 0.4
        max_mask_samples = int(sr * max_mask_time)
        
        # 根据音频长度决定遮盖次数：
        # 比如每 3 秒钟允许遮盖 1 次
        num_masks = max(1, audio_len // (sr * 3)) 
        
        for _ in range(num_masks):
            # 每次随机遮盖 0.1s 到 0.4s
            current_mask_len = random.randint(int(sr * 0.1), max_mask_samples)
            
            if audio_len > current_mask_len:
                start_pos = random.randint(0, audio_len - current_mask_len)
                
                # 填充微小噪音（模拟环境底噪）
                noise = torch.randn(1, current_mask_len).to(waveform.device) * 0.002
                waveform[:, start_pos : start_pos + current_mask_len] = noise
                
        return waveform
    
    def __getitem__(self, index) -> BatchItem:
        row: TsvFormat = self.data.iloc[index]
        audio_path: Path = self.audio_dir / row['path']
        # text: str = unicodedata.normalize('NFC', row['sentence'].strip()) 
        text: str = normalize_extended_uyghur_characters(collapse_spaces(row['sentence'].strip()))

        waveform = self._load_audio(audio_path=audio_path)
        waveform = self._augment_waveform(waveform)
        waveform = waveform.squeeze(0)

        return BatchItem(
            waveform=waveform,
            target_ids=torch.tensor(self.tokenizer.encode(text), dtype=torch.long),
            target_text=text,
            audio_path=str(audio_path)
        )
    
def collate_fn(items: List[BatchItem]) -> Batch:
    max_waveform_len = max(item['waveform'].shape[0] for item in items)
    max_target_len = max(len(item['target_ids']) for item in items)

    batch_size = len(items)

    waveforms = torch.zeros(batch_size, max_waveform_len)
    targets = torch.zeros(batch_size, max_target_len, dtype=torch.long)
    waveform_lengths = torch.zeros(batch_size, dtype=torch.long)
    target_lengths = torch.zeros(batch_size, dtype=torch.long)

    target_texts = []
    audio_paths = []
    

    for i, item in enumerate(items):
        waveform_len = item['waveform'].shape[0]
        target_len = len(item['target_ids'])

        waveforms[i, :waveform_len] = item['waveform']
        targets[i, :target_len] = item['target_ids']
        waveform_lengths[i] = waveform_len
        target_lengths[i] = target_len

        target_texts.append(item['target_text'])
        audio_paths.append(item['audio_path'])

    return Batch(
        waveforms=waveforms,
        targets=targets,
        waveform_lengths=waveform_lengths,
        target_lengths=target_lengths,
        target_texts=target_texts,
        audio_paths=audio_paths
    )


def create_dataloader(tsv_path: Path, audio_dir: Path, noise_dir: Path, tokenizer: ASRTokenizer, batch_size: int = 8, shuffle: bool = True, augment: bool = True, augment_prob: int = 0.5) -> DataLoader:
    dataset = CommonVoiceDataset(tsv_path=tsv_path, audio_dir=audio_dir, noise_dir=noise_dir, tokenizer=tokenizer, augment=augment, augment_prob=augment_prob)

    return DataLoader(dataset=dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn, pin_memory=True, num_workers=8, prefetch_factor=8, persistent_workers=True)


# ============ 测试代码 ============
if __name__ == "__main__":
    from pathlib import Path
    
    workspace_dir = Path(__file__).parent.parent
    
    # 初始化tokenizer
    tokenizer = ASRTokenizer(workspace_dir / 'config' / 'asr_vocab.json')
    
    # 创建数据加载器
    dataloader = create_dataloader(
        tsv_path=workspace_dir / 'data'  / 'ug' / 'train.tsv',
        audio_dir=workspace_dir / 'data'  / 'ug' / 'clips',
        tokenizer=tokenizer,
        batch_size=2,
        shuffle=True,
    )
    
    # 测试加载一个batch
    print("测试数据加载:")
    for batch in dataloader:
        batch: Batch
        print(f"Mel specs shape: {batch['mel_specs'].shape}")
        print(f"Targets shape: {batch['targets'].shape}")
        print(f"Mel lengths: {batch['mel_lengths']}")
        print(f"Target lengths: {batch['target_lengths']}")
        print(f"Target texts: {batch['target_texts']}")
        print(f"Audio paths: {batch['audio_paths']}")
        break





import torch
import torchaudio
from silero_vad import load_silero_vad, read_audio, get_speech_timestamps

# --- 1. 加载 VAD 模型 ---
# 推荐使用 GPU (cuda) 如果可用，否则使用 CPU (cpu)
vad_model = load_silero_vad(source="local", force_onnx=False)

def process_audio_with_vad_and_asr(audio_path, inference_module):
    """
    使用 Silero VAD 分割音频，然后对每个语音片段进行 ASR 转录。

    Args:
        audio_path (str): 输入音频文件的路径。
        inference_module: 您封装了 transcribe 方法的模块或对象。
                          需要确保其 _load_audio 和 transcribe 方法可用。
    """
    print(f"正在加载音频: {audio_path}")
    
    # --- 2. 加载音频用于 VAD 检测 ---
    # Silero VAD 推荐使用 16kHz 采样率
    waveform_vad, sample_rate_vad = torchaudio.load(audio_path)
    if sample_rate_vad != 16000:
        # 如果采样率不是16kHz，需要重采样以供VAD使用
        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate_vad, new_freq=16000)
        waveform_vad = resampler(waveform_vad)
        sample_rate_vad = 16000

    # --- 3. 获取语音时间段 ---
    # get_speech_timestamps 返回一个列表，每个元素是 {'start': start_sample, 'end': end_sample} 的字典
    # 采样率是16000，所以时间戳单位是 1/16000 秒
    speech_timestamps = get_speech_timestamps(
        waveform_vad, 
        vad_model, 
        sampling_rate=sample_rate_vad,
        threshold=0.5,  # 可以根据需要调整阈值
        min_speech_duration_ms=250, # 最小语音持续时间，防止短噪音被误判
        max_speech_duration_s=float('inf'), # 最大语音持续时间，float('inf') 表示不限制
        min_silence_duration_ms=100, # 最小静音间隔，用于分割语音块
        window_size_samples=1536, # VAD窗口大小
        speech_pad_ms=30 # 在语音块前后添加的填充时间
    )

    print(f"检测到 {len(speech_timestamps)} 个语音片段。")

    full_text = ""
    for i, ts in enumerate(speech_timestamps):
        start_sample = int(ts['start'])
        end_sample = int(ts['end'])

        # 计算时间戳（秒）
        start_time = start_sample / sample_rate_vad
        end_time = end_sample / sample_rate_vad
        
        print(f"\n处理第 {i+1} 个片段: 时间范围 [{start_time:.2f}s - {end_time:.2f}s]")

        # --- 4. 从原始音频中提取此片段 ---
        # 注意：这里假设您的 transcribe 函数可以接受 waveform 张量。
        # 我们需要从原始可能不同采样率的音频中提取片段，或者用VAD处理过的waveform。
        # 为了匹配您原来的 _load_audio 方式，我们用 torchaudio 再次精确加载片段。
        
        # 计算原始音频中的样本索引（如果原始音频采样率与VAD不同）
        original_waveform, original_sr = torchaudio.load(audio_path)
        if sample_rate_vad != original_sr:
             # 如果VAD和原始音频采样率不同，需要重新映射索引
            start_idx_orig = int(start_sample * (original_sr / sample_rate_vad))
            end_idx_orig = int(end_sample * (original_sr / sample_rate_vad))
        else:
            start_idx_orig = start_sample
            end_idx_orig = end_sample
            
        segment_waveform = original_waveform[:, start_idx_orig:end_idx_orig]

        # --- 5. 对该片段进行 ASR 转录 ---
        # 这里调用您原有的 transcribe 方法
        try:
            # 假设 transcribe 方法接受一个 waveform tensor
            segment_text = inference_module.transcribe(waveform=segment_waveform)
            
            print(f"  -> 转录结果: {segment_text}")
            full_text += f"[{start_time:.2f}-{end_time:.2f}s] {segment_text}\n"
        except Exception as e:
            print(f"  -> 转录第 {i+1} 个片段时出错: {e}")

    print("\n--- 完整转录结果 ---")
    print(full_text)


# --- 使用示例 ---
# 假设您有一个名为 'inference' 的模块对象，它有 _load_audio 和 transcribe 方法
# audio_path = "your_audio_file.wav"
# process_audio_with_vad_and_asr(audio_path, inference)