import re
import sys
import yaml
import numpy as np
import librosa
import torch
import phonemizer
import noisereduce as nr
from munch import Munch

from meldataset import TextCleaner
from models import ProsodyPredictor, TextEncoder, StyleEncoder
from Modules.hifigan import Decoder


# -------------------------
# Windows-only espeak-ng loader
# -------------------------
if sys.platform.startswith("win"):
    try:
        from phonemizer.backend.espeak.wrapper import EspeakWrapper
        import espeakng_loader
        EspeakWrapper.set_library(espeakng_loader.get_library_path())
    except Exception as e:
        print(e)

_TOKEN_RE = re.compile(r"\S+")

def normalize_phonem_tokens(phonem: str) -> str:
    return " ".join(_TOKEN_RE.findall((phonem or "").strip()))

def espeak_phn(text: str, lang: str) -> str:
    """
    Nếu phonemizer/espeak lỗi -> raise để bạn biết ngay thiếu espeak-ng / libespeak-ng1 / voice 'vi'
    """
    try:
        backend = phonemizer.backend.EspeakBackend(
            language=lang,
            preserve_punctuation=True,
            with_stress=True,
            language_switch="remove-flags",
        )
        out = backend.phonemize([text])[0]
        out = (out or "").strip()
        if len(out) == 0:
            raise RuntimeError(f"phonemizer returned empty output for lang='{lang}', text='{text[:50]}'")
        return out
    except Exception as e:
        raise RuntimeError(f"espeak/phonemizer failed (lang={lang}). Error: {e}")

class Preprocess:
    def __text_normalize(self, text):
        punctuation = ["，", "、", "،", ";", "(", "．", "。", "…", "!", "–", ":", "?"]
        map_to = "."
        punctuation_pattern = re.compile(f"[{''.join(re.escape(p) for p in punctuation)}]")
        text = punctuation_pattern.sub(map_to, text)
        text = re.sub(r"\s+", " ", text).strip()
        return text

    def __merge_fragments(self, texts, n):
        merged = []
        i = 0
        while i < len(texts):
            fragment = texts[i]
            j = i + 1
            while len(fragment.split()) < n and j < len(texts):
                fragment += ", " + texts[j]
                j += 1
            merged.append(fragment)
            i = j

        if len(merged) > 1 and len(merged[-1].split()) < n:
            merged[-2] = merged[-2] + ", " + merged[-1]
            del merged[-1]
        return merged

    def wave_preprocess(self, wave, sr=24000):
        wave = np.asarray(wave, dtype=np.float32).squeeze()
        mel = librosa.feature.melspectrogram(
            y=wave,
            sr=sr,
            n_fft=2048,
            win_length=1200,
            hop_length=300,
            n_mels=80,
            power=2.0,
        )  # (80, T)

        mean, std = -4, 4
        mel = np.log(1e-5 + mel)
        mel = (mel - mean) / std
        return torch.from_numpy(mel).float().unsqueeze(0)  # (1, 80, T)

    def text_preprocess(self, text, n_merge=12):
        text_norm = self.__text_normalize(text).split(".")
        text_norm = [s.strip() for s in text_norm if s.strip()]
        return self.__merge_fragments(text_norm, n=n_merge)

    def length_to_mask(self, lengths):
        mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
        return torch.gt(mask + 1, lengths.unsqueeze(1))

class StyleTTS2(torch.nn.Module):
    def __init__(self, config_path, models_path):
        super().__init__()
        self.register_buffer("get_device", torch.empty(0))
        self.preprocess = Preprocess()

        config = yaml.safe_load(open(config_path, "r", encoding="utf-8"))

        symbols = (
            list(config["symbol"]["pad"])
            + list(config["symbol"]["punctuation"])
            + list(config["symbol"]["letters"])
            + list(config["symbol"]["letters_ipa"])
            + list(config["symbol"]["extend"])
        )
        symbol_dict = {s: i for i, s in enumerate(symbols)}
        n_token = len(symbol_dict) + 1
        print("\nFound:", n_token, "symbols")

        args = self.__recursive_munch(config["model_params"])
        args["n_token"] = n_token

        self.cleaner = TextCleaner(symbol_dict, debug=True)

        self.decoder = Decoder(
            dim_in=args.hidden_dim,
            style_dim=args.style_dim,
            dim_out=args.n_mels,
            resblock_kernel_sizes=args.decoder.resblock_kernel_sizes,
            upsample_rates=args.decoder.upsample_rates,
            upsample_initial_channel=args.decoder.upsample_initial_channel,
            resblock_dilation_sizes=args.decoder.resblock_dilation_sizes,
            upsample_kernel_sizes=args.decoder.upsample_kernel_sizes,
        )
        self.predictor = ProsodyPredictor(
            style_dim=args.style_dim,
            d_hid=args.hidden_dim,
            nlayers=args.n_layer,
            max_dur=args.max_dur,
            dropout=args.dropout,
        )
        self.text_encoder = TextEncoder(
            channels=args.hidden_dim,
            kernel_size=5,
            depth=args.n_layer,
            n_symbols=args.n_token,
        )
        self.style_encoder = StyleEncoder(
            dim_in=args.dim_in,
            style_dim=args.style_dim,
            max_conv_dim=args.hidden_dim,
        )

        n_speakers = config["data_params"]["n_speakers"]
        
        self.spk_emb = torch.nn.Embedding(
            n_speakers,
            args.style_dim
        )
        
        self.spk_ln = torch.nn.LayerNorm(args.style_dim)


        self.__load_models(models_path)

    def text_to_sequence_char_level(text, symbol_dict):
        seq = []
        for ch in text:
            if ch == " ":
                continue
            if ch in symbol_dict:
                seq.append(symbol_dict[ch])
            else:
                print("[WARN] dropped char:", repr(ch))
        return seq


    def __recursive_munch(self, d):
        if isinstance(d, dict):
            return Munch((k, self.__recursive_munch(v)) for k, v in d.items())
        if isinstance(d, list):
            return [self.__recursive_munch(v) for v in d]
        return d

    def __replace_outliers_zscore(self, tensor, threshold=3.0, factor=0.95):
        mean = tensor.mean()
        std = tensor.std()
        z = (tensor - mean) / (std + 1e-8)
        outlier_mask = torch.abs(z) > threshold
        sign = torch.sign(tensor - mean)
        replacement = mean + sign * (threshold * std * factor)
        result = tensor.clone()
        result[outlier_mask] = replacement[outlier_mask]
        return result

    def __load_models(self, models_path):
        # model = {
        #     "decoder": self.decoder,
        #     "predictor": self.predictor,
        #     "text_encoder": self.text_encoder,
        #     "style_encoder": self.style_encoder,
        # }

        model = {
            "decoder": self.decoder,
            "predictor": self.predictor,
            "text_encoder": self.text_encoder,
            "style_encoder": self.style_encoder,
            "spk_emb": self.spk_emb,
            "spk_ln": self.spk_ln,
        }


        params_whole = torch.load(models_path, map_location="cpu")
        params = params_whole["net"]
        params = {k: v for k, v in params.items() if k in model}

        for k in model:
            try:
                model[k].load_state_dict(params[k])
            except Exception:
                from collections import OrderedDict
                new_state_dict = OrderedDict()
                for kk, vv in params[k].items():
                    new_state_dict[kk[7:]] = vv  # strip "module."
                model[k].load_state_dict(new_state_dict, strict=False)

            print(k, ":", sum(p.numel() for p in model[k].parameters()))

    def __compute_style(self, path, denoise, split_dur):
        device = self.get_device.device
        denoise = min(float(denoise), 1.0)
        split_dur = int(split_dur) if split_dur else 0

        wave, sr = librosa.load(path, sr=24000)
        audio, _ = librosa.effects.trim(wave, top_db=30)

        if denoise > 0.0:
            audio_denoise = nr.reduce_noise(
                y=audio, sr=sr, n_fft=2048, win_length=1200, hop_length=300
            )
            audio = audio * (1 - denoise) + audio_denoise * denoise

        with torch.no_grad():
            if split_dur > 0 and len(audio) / sr >= 4:
                jump = sr * split_dur
                total_len = len(audio)
                ref_s = None
                count = 0

                for i in range(0, total_len, jump):
                    seg = audio[i : min(i + jump, total_len)]
                    if len(seg) < sr:  # <1s thì bỏ
                        continue
                    mel = self.preprocess.wave_preprocess(seg).to(device)
                    s = self.style_encoder(mel.unsqueeze(1))
                    ref_s = s if ref_s is None else (ref_s + s)
                    count += 1

                if ref_s is None:
                    mel = self.preprocess.wave_preprocess(audio).to(device)
                    ref_s = self.style_encoder(mel.unsqueeze(1))
                else:
                    ref_s = ref_s / count
            else:
                mel = self.preprocess.wave_preprocess(audio).to(device)
                ref_s = self.style_encoder(mel.unsqueeze(1))

        return ref_s

    def __inference(self, phonem, ref_s, speed=1.0, prev_d_mean=0.0, t=0.1):

        device = self.get_device.device
        tokens = self.cleaner(phonem)
        print("[DBG] token_len =", len(tokens))
        print("[DBG] phonem_head =", phonem[:80])
        
        if len(tokens) == 0:
            raise RuntimeError("Token sequence is empty!")
        
        tokens = [0] + tokens + [0]
        tokens = torch.LongTensor(tokens).unsqueeze(0).to(device)
              
        print("\n========== TOKEN DEBUG ==========")
        print("Max token id:", tokens.max().item())
        print("Min token id:", tokens.min().item())
        print("First 50 tokens:", tokens[0][:50].tolist())
        print("=================================\n")
        

        with torch.no_grad():
            input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
            text_mask = self.preprocess.length_to_mask(input_lengths).to(device)

            t_en = self.text_encoder(tokens, input_lengths, text_mask)
            s = ref_s.to(device)
            if hasattr(self, "spk_emb"):
                spk_id = torch.LongTensor([self.current_speaker_id]).to(device)
                spk_vec = self.spk_emb(spk_id)          # [1, style_dim]
                spk_vec = self.spk_ln(spk_vec)
            
                # merge speaker embedding into style
                s = s + spk_vec
            
                print("🎤 Speaker embedding injected:", self.current_speaker_id)
            
            print("🎤 Style vec mean/std:", s.mean().item(), s.std().item())

            d = self.predictor.text_encoder(t_en, s, input_lengths, text_mask)
            x, _ = self.predictor.lstm(d)
            duration = self.predictor.duration_proj(x)
            duration = torch.sigmoid(duration).sum(dim=-1)

            if prev_d_mean != 0:
                dur_stats = torch.empty_like(duration).normal_(mean=prev_d_mean, std=duration.std() + 1e-8).to(device)
            else:
                dur_stats = torch.empty_like(duration).normal_(mean=duration.mean(), std=duration.std() + 1e-8).to(device)

            duration = duration * (1 - t) + dur_stats * t
            duration[:, 1:-2] = self.__replace_outliers_zscore(duration[:, 1:-2])
            duration = duration / speed

            pred_dur = torch.round(duration.squeeze(0)).clamp(min=1)

            L = int(input_lengths.item())
            T = int(pred_dur.sum().item())
            pred_aln_trg = torch.zeros((L, T), device=device)

            c = 0
            for i in range(L):
                di = int(pred_dur[i].item())
                pred_aln_trg[i, c : c + di] = 1
                c += di

            alignment = pred_aln_trg.unsqueeze(0)

            en = d.transpose(-1, -2) @ alignment
            F0_pred, N_pred = self.predictor.F0Ntrain(en, s)
            asr = t_en @ pred_aln_trg.unsqueeze(0)

            out = self.decoder(asr, F0_pred, N_pred, s)

        return out.squeeze().cpu().numpy(), float(duration.mean().item())

    def get_styles(self, speakers, denoise=0.3, avg_style=True):
        split_dur = 2 if avg_style else 0
        styles = {}
        for sid, meta in speakers.items():
            ref_s = self.__compute_style(meta["path"], denoise=denoise, split_dur=split_dur)
            styles[sid] = {
                "style": ref_s,
                "path": meta["path"],
                "lang": meta["lang"],
                "speed": meta["speed"],
            }
        return styles
    def generate(self, text, styles, stabilize=True, n_merge=16, default_speaker="[id_1]"):
    
        smooth_value = 0.2 if stabilize else 0.0
    
        list_wav = []
        prev_d_mean = 0.0
        text = re.sub(r"[\n\r\t\f\v]", "", text)
    
        # split by speaker tags
        parts = re.split(r"(\[id_\d+\])", text)
        if len(parts) <= 1 or re.match(r"(\[id_\d+\])", parts[0]) is None:
            parts.insert(0, default_speaker)
    
        speaker_tag = None        # "id_1"
        speaker_id = None         # int
        current_ref_s = None
        speed = 1.0
    
        for p in parts:
            # -----------------------------
            # Parse speaker tag
            # -----------------------------
            if re.match(r"(\[id_\d+\])", p):
                speaker_tag = p.strip("[]")                     # "id_1"
                speaker_id = int(speaker_tag.replace("id_", ""))
    
                # expose speaker id for inference
                self.current_speaker_id = speaker_id
    
                if speaker_tag not in styles:
                    raise RuntimeError(f"Speaker {speaker_tag} not found in styles!")
    
                current_ref_s = styles[speaker_tag]["style"]
                speed = styles[speaker_tag]["speed"]
    
                print("🎤 Active speaker:", speaker_tag, "->", speaker_id)
                continue
    
            if not p.strip():
                continue
    
            if speaker_id is None:
                raise RuntimeError("Speaker ID chưa được set! Kiểm tra tag [id_x].")
    
            # -----------------------------
            # Text → phoneme → waveform
            # -----------------------------
            for sentence in self.preprocess.text_preprocess(p, n_merge=n_merge):
                phonem = espeak_phn(sentence, styles[speaker_tag]["lang"])
    
                wav, prev_d_mean = self.__inference(
                    phonem,
                    current_ref_s,
                    speed=speed,
                    prev_d_mean=prev_d_mean,
                    t=smooth_value,
                )
    
                # -----------------------------
                # Debug
                # -----------------------------
                print("[DBG] wav shape:", wav.shape)
                print("[DBG] wav min/max:", wav.min().item(), wav.max().item())
                print("[DBG] wav mean abs:", np.abs(wav).mean())
    
                # -----------------------------
                # Safe trim
                # -----------------------------
                trim = int(0.05 * 24000)   # 50 ms
                if wav.shape[0] > 4 * trim:
                    wav = wav[trim:-trim]
    
                if wav.size > 0:
                    list_wav.append(wav)
    
        # -----------------------------
        # Merge all chunks
        # -----------------------------
        if len(list_wav) == 0:
            print("⚠️ No audio generated → return silence")
            return np.zeros((2400,), dtype=np.float32)
    
        final_wav = np.concatenate(list_wav)
    
        # pad head & tail
        pad = int(0.05 * 24000)
        final_wav = np.concatenate(
            [np.zeros((pad,), dtype=np.float32), final_wav, np.zeros((pad,), dtype=np.float32)]
        )
    
        return final_wav