from __future__ import annotations

import argparse
import contextlib
import io
import logging
import re
import sys
import warnings
from pathlib import Path

import numpy as np
import torch
from safetensors.numpy import load_file

from configuration import cleaned_text_to_sequence, get_hparams_from_file, symbols
from modeling import SynthesizerTrn

PACKAGE_ROOT = Path(__file__).resolve().parent


def split_text(text: str, limit: int = 280) -> list[str]:
    normalized = " ".join(text.split())
    sentences = [
        part.strip()
        for part in re.split(r"(?<=[.!?;:])\s+", normalized)
        if part.strip()
    ]
    chunks: list[str] = []
    for sentence in sentences or [normalized]:
        while len(sentence) > limit:
            search = sentence[: limit + 1]
            punctuation = max(search.rfind(mark) for mark in (",", ";", ":"))
            split_at = (
                punctuation + 1
                if punctuation >= limit // 2
                else sentence.rfind(" ", 0, limit + 1)
            )
            if split_at < limit // 2:
                split_at = limit
            chunks.append(sentence[:split_at].strip())
            sentence = sentence[split_at:].strip()
        if sentence:
            chunks.append(sentence)
    return chunks


def boundary_pause_seconds(chunk: str) -> float:
    ending = chunk.rstrip()[-1:] if chunk.strip() else ""
    return {
        "?": 0.28,
        "!": 0.24,
        ".": 0.22,
        ";": 0.16,
        ":": 0.13,
        ",": 0.09,
    }.get(ending, 0.08)


def edge_fade(waveform: np.ndarray, sample_rate: int, milliseconds: float = 5.0) -> np.ndarray:
    frames = min(round(sample_rate * milliseconds / 1000.0), waveform.size // 2)
    if frames <= 0:
        return waveform
    output = waveform.copy()
    ramp = np.linspace(0.0, 1.0, frames, endpoint=True, dtype=np.float32)
    output[:frames] *= ramp
    output[-frames:] *= ramp[::-1]
    return output


def optimize_for_inference(model: SynthesizerTrn) -> None:
    """Collapse weight normalization for VTXVocoder & VTXVectorEstimator."""
    with contextlib.redirect_stdout(io.StringIO()):
        model.VTXVocoder.remove_weight_norm()
    for flow in model.VTXVectorEstimator.flows:
        encoder = getattr(flow, "enc", None)
        if encoder is not None and hasattr(encoder, "remove_weight_norm"):
            encoder.remove_weight_norm()


class VtxTTS:
    """VTX-TTS Engine using native 4-Bit model.safetensors."""

    def __init__(self, model_dir: str | Path = PACKAGE_ROOT, device: str = "cpu") -> None:
        self.root = Path(model_dir).resolve()
        self.device = torch.device(device)
        self.hps = get_hparams_from_file(str(self.root / "config.json"))
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message="`torch.nn.utils.weight_norm` is deprecated",
                category=FutureWarning,
            )
            model_kwargs = self.hps.model.__dict__ if hasattr(self.hps.model, "__dict__") else self.hps.model
            self.model = SynthesizerTrn(
                len(symbols),
                self.hps.data.filter_length // 2 + 1,
                self.hps.train.segment_size // self.hps.data.hop_length,
                **model_kwargs,
            ).to(self.device).eval()

        # Load native 4-bit safetensors
        st_file = self.root / "model.safetensors"
        st_weights = load_file(str(st_file))
        
        state_dict = {}
        for k, v in st_weights.items():
            if k.endswith(".packed"):
                base_k = k[:-7]
                scales = st_weights[f"{base_k}.scales"]
                zeros = st_weights[f"{base_k}.zeros"]
                shape = tuple(st_weights[f"{base_k}.shape"])
                
                # Dynamic 4-bit dequantization into model parameter buffer
                low = (v & 0x0F).astype(np.float32)
                high = ((v >> 4) & 0x0F).astype(np.float32)
                unpacked = np.empty((v.size * 2,), dtype=np.float32)
                unpacked[0::2] = low
                unpacked[1::2] = high
                
                n_orig = np.prod(shape)
                blocked = unpacked[:((n_orig + 31)//32)*32].reshape(-1, 32)
                s = scales.astype(np.float32)
                z = zeros.astype(np.float32)
                deq_w = (blocked * s + z).reshape(-1)[:n_orig].reshape(shape)
                state_dict[base_k] = torch.from_numpy(deq_w).float()
            elif k.endswith(".scales") or k.endswith(".zeros") or k.endswith(".shape"):
                continue
            else:
                state_dict[k] = torch.from_numpy(v).float()

        self.model.load_state_dict(state_dict, strict=False)
        optimize_for_inference(self.model)
        self.sample_rate = int(self.hps.data.sampling_rate)

        if self.device.type == "cpu":
            if hasattr(torch, "set_num_threads") and torch.get_num_threads() > 4:
                torch.set_num_threads(4)

    @staticmethod
    def _quantize_lf4(tensor: torch.Tensor, group_size: int = 32) -> torch.Tensor:
        if tensor.numel() < group_size:
            return tensor
        shape = tensor.shape
        flat = tensor.reshape(-1)
        n = flat.numel()
        pad = (group_size - (n % group_size)) % group_size
        if pad > 0:
            flat = torch.cat([flat, torch.zeros(pad, device=tensor.device)])
        
        groups = flat.reshape(-1, group_size)
        g_max = groups.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-8)
        scales = g_max / 7.0
        
        q_groups = torch.round(groups / scales).clamp(-7, 7)
        deq_groups = q_groups * scales
        return deq_groups.reshape(-1)[:n].reshape(shape)

    def _tokens(self, text: str) -> tuple[torch.Tensor, torch.Tensor]:
        from phonemizer.backend import EspeakBackend
        backend = EspeakBackend('en-us', preserve_punctuation=True, with_stress=True)
        phoneme_str = backend.phonemize([text])[0]
        sequence = cleaned_text_to_sequence(phoneme_str)
        if not sequence:
            sequence = [1, 2, 3]
        if self.hps.data.add_blank:
            res = [0] * (len(sequence) * 2 + 1)
            res[1::2] = sequence
            sequence = res
        tokens = torch.LongTensor(sequence).to(self.device).unsqueeze(0)
        lengths = torch.LongTensor([tokens.size(1)]).to(self.device)
        return tokens, lengths

    @torch.inference_mode()
    def synthesize(
        self,
        text: str,
        *,
        speed: float = 1.0,
        variation: float = 0.667,
        seed: int = 0,
    ) -> tuple[int, np.ndarray]:
        normalized = " ".join(text.split())
        if not normalized:
            raise ValueError("Text must not be empty.")
        if not 0.5 <= speed <= 2.0:
            raise ValueError("speed must be between 0.5 and 2.0")
        if not 0.0 <= variation <= 1.0:
            raise ValueError("variation must be between 0.0 and 1.0")
        chunks = split_text(normalized)
        pieces: list[np.ndarray] = []
        for index, chunk in enumerate(chunks):
            if index:
                pieces.append(
                    np.zeros(
                        round(self.sample_rate * boundary_pause_seconds(chunks[index - 1])),
                        dtype=np.float32,
                    )
                )
            tokens, lengths = self._tokens(chunk)
            torch.manual_seed(seed + index)
            if self.device.type == "cuda":
                torch.cuda.manual_seed_all(seed + index)
            waveform = self.model.infer(
                tokens,
                lengths,
                noise_scale=variation,
                noise_scale_w=0.8,
                length_scale=1.0 / speed,
                max_len=4000,
            )[0][0, 0].float().cpu().numpy()
            pieces.append(edge_fade(waveform, self.sample_rate))
        waveform = np.clip(np.concatenate(pieces), -1.0, 1.0)
        return self.sample_rate, waveform

    def save(self, text: str, output: str | Path, **kwargs: object) -> Path:
        destination = Path(output)
        destination.parent.mkdir(parents=True, exist_ok=True)
        sample_rate, waveform = self.synthesize(text, **kwargs)
        
        # Zero-dependency native PCM16 WAV writer (eliminates scipy/soundfile requirement)
        data_int16 = (waveform * 32767.0).clip(-32768, 32767).astype("<i2")
        raw_bytes = data_int16.tobytes()
        data_size = len(raw_bytes)
        
        header = bytearray()
        header.extend(b"RIFF")
        header.extend((36 + data_size).to_bytes(4, "little"))
        header.extend(b"WAVE")
        header.extend(b"fmt ")
        header.extend((16).to_bytes(4, "little"))
        header.extend((1).to_bytes(2, "little"))
        header.extend((1).to_bytes(2, "little"))
        header.extend((sample_rate).to_bytes(4, "little"))
        header.extend((sample_rate * 2).to_bytes(4, "little"))
        header.extend((2).to_bytes(2, "little"))
        header.extend((16).to_bytes(2, "little"))
        header.extend(b"data")
        header.extend((data_size).to_bytes(4, "little"))
        
        with open(destination, "wb") as f:
            f.write(header)
            f.write(raw_bytes)
            
        return destination


def main() -> None:
    parser = argparse.ArgumentParser(description="Run standalone VTX-TTS speech synthesis.")
    parser.add_argument("--model-dir", type=Path, default=PACKAGE_ROOT)
    parser.add_argument("--text", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--speed", type=float, default=1.0)
    parser.add_argument("--variation", type=float, default=0.667)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()
    engine = VtxTTS(args.model_dir, args.device)
    engine.save(
        args.text,
        args.output,
        speed=args.speed,
        variation=args.variation,
        seed=args.seed,
    )
    print(f"wrote {args.output} at {engine.sample_rate} Hz")


if __name__ == "__main__":
    main()