import io
import sys
import types

import torch
import torchaudio
import s3tokenizer
import onnxruntime
import numpy as np

import torchaudio.compliance.kaldi as kaldi
from stepaudio2.flashcosyvoice.modules.hifigan import HiFTGenerator
from stepaudio2.flashcosyvoice.utils.audio import mel_spectrogram
from hyperpyyaml import load_hyperpyyaml


def _setup_cosyvoice2_alias():
    """给 hyperpyyaml 提供 cosyvoice2.* 的兼容别名，不改 flow.yaml。"""
    if 'cosyvoice2.flow.flow' in sys.modules:
        # 已经设置过，直接复用
        return

    # 导入 stepaudio2 里真实的实现
    import stepaudio2.cosyvoice2.flow.flow as _step_flow
    import stepaudio2.cosyvoice2.flow.flow_matching as _step_flow_matching
    import stepaudio2.cosyvoice2.flow.decoder_dit as _step_decoder_dit
    import stepaudio2.cosyvoice2.transformer.upsample_encoder_v2 as _step_upsample

    # 创建顶层 cosyvoice2 包和子包
    cosyvoice2_pkg = types.ModuleType('cosyvoice2')
    cosyvoice2_flow_pkg = types.ModuleType('cosyvoice2.flow')
    cosyvoice2_transformer_pkg = types.ModuleType('cosyvoice2.transformer')

    # 挂载子模块
    cosyvoice2_flow_pkg.flow = _step_flow
    cosyvoice2_flow_pkg.flow_matching = _step_flow_matching
    cosyvoice2_flow_pkg.decoder_dit = _step_decoder_dit

    cosyvoice2_transformer_pkg.upsample_encoder_v2 = _step_upsample

    cosyvoice2_pkg.flow = cosyvoice2_flow_pkg
    cosyvoice2_pkg.transformer = cosyvoice2_transformer_pkg

    # 注册到 sys.modules，让 `cosyvoice2.flow.flow.*` 这类路径可被 import
    sys.modules['cosyvoice2'] = cosyvoice2_pkg
    sys.modules['cosyvoice2.flow'] = cosyvoice2_flow_pkg
    sys.modules['cosyvoice2.flow.flow'] = _step_flow
    sys.modules['cosyvoice2.flow.flow_matching'] = _step_flow_matching
    sys.modules['cosyvoice2.flow.decoder_dit'] = _step_decoder_dit
    sys.modules['cosyvoice2.transformer'] = cosyvoice2_transformer_pkg
    sys.modules['cosyvoice2.transformer.upsample_encoder_v2'] = _step_upsample

def fade_in_out(fade_in_mel:torch.Tensor, fade_out_mel:torch.Tensor, window:torch.Tensor):
    """perform fade_in_out in tensor style
    """
    mel_overlap_len = int(window.shape[0] / 2)
    fade_in_mel = fade_in_mel.clone()
    fade_in_mel[..., :mel_overlap_len] = \
        fade_in_mel[..., :mel_overlap_len] * window[:mel_overlap_len] + \
        fade_out_mel[..., -mel_overlap_len:] * window[mel_overlap_len:]
    return fade_in_mel


class Token2wav():
    def __init__(self, model_path, float16=False, n_timesteps=10):
        self.float16 = float16
        self.n_timesteps = n_timesteps

        # 在加载 flow.yaml 之前，先把 cosyvoice2 的别名注册好
        _setup_cosyvoice2_alias()

        self.audio_tokenizer = s3tokenizer.load_model(f"{model_path}/speech_tokenizer_v2_25hz.onnx").cuda().eval()

        option = onnxruntime.SessionOptions()
        option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
        option.intra_op_num_threads = 1
        self.spk_model = onnxruntime.InferenceSession(f"{model_path}/campplus.onnx", sess_options=option, providers=["CPUExecutionProvider"])

        with open(f"{model_path}/flow.yaml", "r") as f:
            configs = load_hyperpyyaml(f)
            self.flow = configs['flow']
        if float16:
            self.flow.half()
        self.flow.load_state_dict(torch.load(f"{model_path}/flow.pt", map_location="cpu", weights_only=True), strict=True)
        self.flow.cuda().eval()

        self.hift = HiFTGenerator()
        hift_state_dict = {k.replace('generator.', ''): v for k, v in torch.load(f"{model_path}/hift.pt", map_location="cpu", weights_only=True).items()}
        self.hift.load_state_dict(hift_state_dict, strict=True)
        self.hift.cuda().eval()

        self.cache = None

        # stream conf
        self.mel_cache_len = 8  # hard-coded, 160ms
        self.source_cache_len = int(self.mel_cache_len * 480)   # 50hz mel -> 24kHz wave
        self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).cuda()

        # hifigan cache
        self.hift_cache_dict = {}

    def _prepare_prompt(self, prompt_wav):
        audio = s3tokenizer.load_audio(prompt_wav, sr=16000)  # [T]
        # TODO 在audio 后面 pad 3个token长度 = 0.12s的音频
        mels = s3tokenizer.log_mel_spectrogram(audio)
        mels, mels_lens = s3tokenizer.padding([mels])
        prompt_speech_tokens, prompt_speech_tokens_lens = self.audio_tokenizer.quantize(mels.cuda(), mels_lens.cuda())
        # TODO 手动赋值最后3个token为 4218 静音token

        spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000)
        spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True)
        spk_emb = torch.tensor(self.spk_model.run(
            None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()}
        )[0], device='cuda')

        audio, sample_rate = torchaudio.load(prompt_wav, backend='soundfile')
        audio = audio.mean(dim=0, keepdim=True)  # [1, T]
        if sample_rate != 24000:
            audio = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=24000)(audio)
        prompt_mel = mel_spectrogram(audio).transpose(1, 2).squeeze(0)  # [T, num_mels]
        prompt_mels = prompt_mel.unsqueeze(0).cuda()
        prompt_mels_lens = torch.tensor([prompt_mels.shape[1]], dtype=torch.int32, device='cuda')
        prompt_mels = torch.nn.functional.pad(prompt_mels, (0, 0, 0, prompt_speech_tokens.shape[1] * self.flow.up_rate - prompt_mels.shape[1]), mode='replicate')
        return prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens

    def __call__(self, generated_speech_tokens, prompt_wav):
        if self.cache is None:
            self.cache = self._prepare_prompt(prompt_wav)
        prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = self.cache

        generated_speech_tokens = torch.tensor([generated_speech_tokens], dtype=torch.int32, device='cuda')
        generated_speech_tokens_lens = torch.tensor([generated_speech_tokens.shape[1]], dtype=torch.int32, device='cuda')

        with torch.amp.autocast("cuda", dtype=torch.float16 if self.float16 else torch.float32):
            mel = self.flow.inference(generated_speech_tokens, generated_speech_tokens_lens,
                prompt_speech_tokens, prompt_speech_tokens_lens,
                prompt_mels, prompt_mels_lens, spk_emb, 10)

        wav, _ = self.hift(speech_feat=mel)
        output = io.BytesIO()
        torchaudio.save(output, wav.cpu(), sample_rate=24000, format='wav')

        return output.getvalue()

    def set_stream_cache(self, prompt_wav):
        self.cache = self._prepare_prompt(prompt_wav)
        prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = self.cache

        print("prompt_speech_tokens_lens", prompt_speech_tokens_lens)
        print("prompt_mels.shape", prompt_mels.shape)
        print("prompt_mels_lens", prompt_mels_lens)
        
        right_pad_speech_tokens = torch.ones(1, 3, device=prompt_speech_tokens.device, dtype=prompt_speech_tokens.dtype) * 4218
        
        # self.stream_cache = self.flow.setup_cache(
        #     torch.cat([prompt_speech_tokens, prompt_speech_tokens[:, :3]], dim=1),
        #     prompt_mels, spk_emb, n_timesteps=self.n_timesteps)
        
        stream_cache = self.flow.setup_cache(
            torch.cat([prompt_speech_tokens, right_pad_speech_tokens], dim=1),
            prompt_mels, spk_emb, n_timesteps=self.n_timesteps)

        # hift cache
        hift_cache_dict = dict(
            mel = torch.zeros(1, prompt_mels.shape[2], 0, device='cuda'), 
            source = torch.zeros(1, 1, 0, device='cuda'),
            speech = torch.zeros(1, 0, device='cuda'),
        )

        return stream_cache, hift_cache_dict


    def stream(self, generated_speech_tokens, prompt_wav, last_chunk=False, return_waveform=False):
        if self.cache is None:
            self.cache = self._prepare_prompt(prompt_wav)
        prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = self.cache

        generated_speech_tokens = torch.tensor([generated_speech_tokens], dtype=torch.int32, device='cuda')
        generated_speech_tokens_lens = torch.tensor([generated_speech_tokens.shape[1]], dtype=torch.int32, device='cuda')

        if self.stream_cache is None:
            raise ValueError("stream_cache is not set")

        with torch.amp.autocast("cuda", dtype=torch.float16 if self.float16 else torch.float32):
            chunk_mel, self.stream_cache = self.flow.inference_chunk(
                token=generated_speech_tokens,
                spk=spk_emb,
                cache=self.stream_cache,
                last_chunk=last_chunk,
                n_timesteps=self.n_timesteps,
            )
        if self.stream_cache['estimator_att_cache'].shape[4] > (prompt_mels.shape[1] + 100):
            self.stream_cache['estimator_att_cache'] = torch.cat([
                self.stream_cache['estimator_att_cache'][:, :, :, :, :prompt_mels.shape[1]],
                self.stream_cache['estimator_att_cache'][:, :, :, :, -100:],
            ], dim=4)

        # bug fix - 20260107
        # 同样截断 conformer_att_cache，防止无限增长导致 position embedding 不匹配
        if self.stream_cache['conformer_att_cache'].shape[3] > (prompt_mels.shape[1] + 100):
            self.stream_cache['conformer_att_cache'] = torch.cat([
                self.stream_cache['conformer_att_cache'][:, :, :, :prompt_mels.shape[1], :],
                self.stream_cache['conformer_att_cache'][:, :, :, -100:, :],
            ], dim=3)

        # vocoder cache
        hift_cache_mel = self.hift_cache_dict['mel']
        hift_cache_source = self.hift_cache_dict['source']
        hift_cache_speech = self.hift_cache_dict['speech']
        mel = torch.concat([hift_cache_mel, chunk_mel], dim=2)

        speech, source = self.hift(mel, hift_cache_source)

        # overlap speech smooth
        if hift_cache_speech.shape[-1] > 0:
            speech = fade_in_out(speech, hift_cache_speech, self.speech_window)

        # 检测是否是第一个 chunk（没有有效的 speech cache）
        is_first_chunk = hift_cache_speech.shape[-1] == 0

        # update vocoder cache
        self.hift_cache_dict = dict(
            mel = mel[..., -self.mel_cache_len:].clone().detach(),
            source = source[:, :, -self.source_cache_len:].clone().detach(),
            speech = speech[:, -self.source_cache_len:].clone().detach(),
        )

        if not last_chunk:
            if is_first_chunk:
                # 第一个 chunk：截断尾部，在开头添加静音来补偿长度
                # 不做淡入，保持原始音频内容不变
                silence_padding = torch.zeros(1, self.source_cache_len, device=speech.device)
                speech = torch.cat([silence_padding, speech[:, :-self.source_cache_len]], dim=1)
            else:
                # 后续 chunk：正常截断尾部（由下一个 chunk 的 fade_in 补偿）
                speech = speech[:, :-self.source_cache_len]

        wav_np = speech.cpu().numpy()
        if return_waveform:
            return wav_np

        # Clip to [-1, 1] to avoid overflow, then scale to int16
        wav_np = np.clip(wav_np, -1.0, 1.0)
        wav_int16 = (wav_np * 32767.0).astype('<i2')  # 16-bit little-endian PCM
        pcm_bytes = wav_int16.tobytes()
        return pcm_bytes

