import os 
import time
import numpy as np
import soundfile as sf
import onnxruntime as ort
from tokenizer import tokenize, chunk_text

def length_regulate_np(text_h, durations):
    """Vectorized duration expansion using np.repeat."""
    B, C, _ = text_h.shape
    mel_lens = durations.sum(axis=1).astype(np.int64)
    T_mel = mel_lens.max()
    
    out = np.zeros((B, C, T_mel), dtype=text_h.dtype)
    for b in range(B):
        expanded = np.repeat(text_h[b], durations[b], axis=-1)
        out[b, :, :expanded.shape[-1]] = expanded[:, :T_mel]
        
    mask = np.arange(T_mel)[None, :] < mel_lens[:, None]
    return out, mel_lens, mask
    
class VachanaFlowTTS():
    def __init__(self, model_path="onnx_model", device="cpu"):
        self.model_path = model_path 
        self.device = device
        
    def load_onnx_all(self, model_path:str, device="cpu"):
        providers = ['CUDAExecutionProvider'] if device == "cuda" else ['CPUExecutionProvider']
        opts = ort.SessionOptions()
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        n_threads = int(os.environ.get("ORT_NUM_THREADS", min(8, os.cpu_count() or 1)))
        opts.intra_op_num_threads = n_threads
        opts.inter_op_num_threads = 1
        encoder_session = ort.InferenceSession(os.path.join(model_path,"tts_encoder.onnx"), providers=providers)
        decoder_session = ort.InferenceSession(os.path.join(model_path,"tts_decoder.onnx"), providers=providers)
        vocoder_session = ort.InferenceSession(os.path.join(model_path,"vocoder.onnx"), providers=providers)
        return encoder_session, decoder_session, vocoder_session

    def infer(self, text:str, spk_id:int = 0, n_steps=8, length_scale=1.15, variation=0.667, normalize=True):
        encoder_session, decoder_session, vocoder_session = self.load_onnx_all(self.model_path, device=self.device)
        token_ids = tokenize(text)
        print(f"text len : {len(text)} | tokens len : {len(token_ids)}")

        text_ids = np.array(token_ids, dtype=np.int64).reshape(1, -1)
        text_lens = np.array([len(token_ids)], dtype=np.int64)
        spk_id = np.array([spk_id], dtype=np.int64)
        
        enc_inputs = {"text_ids": text_ids, "text_lens": text_lens, "spk_ids": spk_id}
        text_h, log_dur, text_mask, spk_emb = encoder_session.run(None, enc_inputs)

        durations = np.exp(log_dur).squeeze(1) * length_scale
        durations = np.clip(np.round(durations), a_min=1.0, a_max=None) * text_mask.astype(np.float32)
        durations = durations.astype(np.int64)

        text_h_exp, mel_lens, mel_mask = length_regulate_np(text_h, durations)
        T_mel = text_h_exp.shape[-1]
        B = text_ids.shape[0]

        x = (np.random.randn(B, 80, T_mel) * variation).astype(np.float32)
        dt = 1.0 / n_steps

        decoder_input_names = [i.name for i in decoder_session.get_inputs()]

        for i in range(n_steps):
            t_val = np.full((B,), i * dt, dtype=np.float32)
            possible_dec_inputs = {"x": x, "t": t_val, "text_h_exp": text_h_exp, "mel_mask": mel_mask, "spk_emb": spk_emb}
            dec_inputs = {k: v for k, v in possible_dec_inputs.items() if k in decoder_input_names}
            v = decoder_session.run(None, dec_inputs)[0]
            x = x + v * dt

        mel_mask_expanded = np.expand_dims(mel_mask, axis=1)
        mel_out_np = np.where(mel_mask_expanded, x, 0.0)
        mel_out_np = mel_out_np[:, :, :mel_lens[0]]

        wav_outputs = vocoder_session.run(None, {"mel": mel_out_np.astype(np.float32)})
        
        audio_batch = wav_outputs[0]
        audio = np.squeeze(audio_batch).astype(np.float32)

        if normalize:
            if np.max(np.abs(audio)) > 0:
                audio = audio / np.max(np.abs(audio))
                
        return audio

    def synthesize(self, text:str, spk_id = 0, n_steps=8, speed=1.0, variation=0.667, normalize=True):
        length_scale = 1.0 / speed
        if len(text) > 1000:
            text_chunk = chunk_text(text)
            audio_chunk = []
            for txt in text_chunk:
                audio = self.infer(txt, spk_id, n_steps, length_scale, variation, normalize)
                audio_chunk.append(audio)
            audio = np.concatenate(audio_chunk)
        else:
            audio = self.infer(text, spk_id, n_steps, length_scale, variation, normalize)
            
        return audio