import torch
import torchaudio
from torch import Tensor, no_grad, device
from torchaudio.transforms import Resample
from pathlib import Path
from typing import TypedDict
from silero_vad import load_silero_vad, get_speech_timestamps
from tokenizer import ASRTokenizer
from model import ASRModel

CONFIG = {
    # 模型配置
    'input_dim': 256,
    'num_heads': 8,
    'ffn_dim': 2048,
    'num_layers': 8,
    'dropout': 0.1,
}

class TimestampsType(TypedDict):
    start: int
    end: int

class ASRInference:
    def __init__(self, model_path: Path, vocab_path: Path, device: device, sample_rate: int = 16000) -> None:
        self.device = device
        self.sample_rate = sample_rate
        self.tokenizer = ASRTokenizer(vocab_path=vocab_path)
        self.model = ASRModel(vocab_size=self.tokenizer.vocab_size(), **CONFIG).to(device=device)
        self.model.load_state_dict(torch.load(model_path, map_location=device)['model_state_dict'])
        self.model.eval()

        print(f"params params: {self.model.get_num_params():,}",)

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

        if sample_rate != self.sample_rate:
            waveform = Resample(sample_rate, self.sample_rate)(waveform)

        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)
        
        waveform = waveform / (waveform.abs().max() + 1e-8)
        return waveform
    
    def transcribe(self, waveform: Tensor) -> str:
        waveform = waveform.to(device=self.device) # [1, time]
        waveform_length = torch.tensor([waveform.shape[1]], dtype=torch.long, device=self.device)

        with no_grad():
            log_probs, _ = self.model(waveforms=waveform, waveform_lengths=waveform_length) # [1, T, vocab]

        text = self.tokenizer.ctc_greedy_decode(log_probs=log_probs[0])
        return text
        
    def transcribe_batch(self, audio_paths: list[Path]) -> list[str]:
        results = []

        for audio_path in audio_paths:
            text = self.transcribe(audio_path=audio_path)
            results.append(text)

        return results
    
    def create(self, audio_path: Path, threshold_ms: int = 20000):
        def merge_timestamps(timestamps: int, threshold_ms: int, sample_rate: int):
            if not timestamps:
                return []
            threshold_samples = (threshold_ms / 1000) * sample_rate
            merged = []
            curr_start = timestamps[0]['start']
            curr_end = timestamps[0]['end']
            
            for i in range(1, len(timestamps)):
                # 如果当前积攒的长度不到 2 秒，就一直合并到当前的 end 上
                if (curr_end - curr_start) < threshold_samples:
                    curr_end = timestamps[i]['end']
                else:
                    merged.append({'start': curr_start, 'end': curr_end})
                    curr_start = timestamps[i]['start']
                    curr_end = timestamps[i]['end']
            merged.append({'start': curr_start, 'end': curr_end})
            return merged

        vad_model = load_silero_vad(onnx=False)
        waveform = self._load_audio(audio_path=audio_path)
        speech_timestamps: list[TimestampsType] = get_speech_timestamps(
            waveform, 
            vad_model, 
            sampling_rate=self.sample_rate,
            threshold=0.5,  # 可以根据需要调整阈值
            min_speech_duration_ms=100, # 最小语音持续时间，防止短噪音被误判
            min_silence_duration_ms=200, # 最小静音间隔，用于分割语音块
            speech_pad_ms=100, # 在语音块前后添加的填充时间
        )
        final_timestamps = merge_timestamps(timestamps=speech_timestamps, threshold_ms=threshold_ms, sample_rate=self.sample_rate)
        for ts in final_timestamps:
            ts: TimestampsType
            start_sample = int(ts['start'])
            end_sample = int(ts['end'])

            segment_waveform = waveform[:, start_sample:end_sample]

            segment_text = self.transcribe(waveform=segment_waveform)
            print(segment_text, end=" ", flush=True)
        print('\n')

def main():
    workspace_dir = Path(__file__).parent.parent
    device = torch.device('cuda:1')

    checkpoint = workspace_dir / ".checkpoints/best_wer_model.pt"
    inference = ASRInference(model_path=checkpoint, vocab_path=workspace_dir / 'config/asr_vocab.json' , device=device)

    audio_path = workspace_dir / 'data/test/F001_001.wav'
    waveform = inference._load_audio(audio_path=audio_path)
    text = inference.transcribe(waveform=waveform)
    print("transcribe:", text)
    inference.create(audio_path=audio_path)

if __name__ == "__main__":
    main()