"""Swar inference: text (Nepali, English or mixed) -> 16 kHz speech.

    from swar import Swar
    tts = Swar.from_pretrained("elitexp/Swar")          # or a local folder
    wav, sr = tts.synthesize("नमस्ते, यहाँलाई स्वागत छ।")  # default voice: bijaya
    wav, sr = tts.synthesize("Your balance is Rs. 250.", voice="path/to/reference.wav",
                             speed=1.0, pause_floor_ms=300, sentence_pause_ms=350)

    python -m swar.inference "नमस्ते, यहाँलाई स्वागत छ।" -o hello.wav

A voice is two 128-d style vectors: acoustic (timbre, read by the decoder) and
prosodic (manner of speaking, read by the duration / pitch / energy predictor).
voices/*.pt hold them ready-made; any clean 3-10 s clip of one speaker gives a new
voice through style_from_audio(). pause_floor_ms makes every comma / full stop /
question mark inside a sentence last at least that long.

Sentences are generated one at a time. Training clips began and ended with room
silence, so the model gives the start and end tokens 0.3-0.5 s each; joined
sentence by sentence that piled up to ~1 s at every full stop and delayed the first
sound. So the edges are cut to EDGE_UNITS, and the gap between two sentences is
exactly sentence_pause_ms, carried by the first sentence's final punctuation.
"""
from __future__ import annotations

import argparse
import re
import time
from pathlib import Path

import numpy as np
import torch

from .mel import SR, Mel
from .loanwords import nepalise
from .model import build, conf
from .text import STOI, normalize, to_ids

PAUSE_CHARS = ",.?!;:…—"
UNIT_S = 0.025                                                     # one duration unit
EDGE_UNITS = 2                                                     # start / end token cap (50 ms)
TAIL_UNITS = 4                                                     # after the last sentence (100 ms)
MAX_CHARS = 220


def _alignment(dur: torch.Tensor, T: int) -> torch.Tensor:
    end = dur.cumsum(1)
    start = end - dur
    t = torch.arange(T, device=dur.device)[None, None]
    return ((t >= start[..., None]) & (t < end[..., None])).float()


def split_sentences(text: str) -> list[str]:
    out = []
    for s in re.split(r"(?<=[.?!])\s+", text):
        s = s.strip()
        while len(s) > MAX_CHARS:                                  # long run-on: cut at a comma or space
            cut = max(s.rfind(",", 0, MAX_CHARS), s.rfind(" ", 0, MAX_CHARS))
            cut = cut if cut > 0 else MAX_CHARS
            out.append(s[: cut + 1].strip())
            s = s[cut + 1:].strip()
        if s:
            out.append(s)
    return out


class Swar:
    def __init__(self, folder: Path, device: str | None = None):
        self.folder = Path(folder)
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.config = conf(self.folder / "config.json")
        m = build("D", self.config)
        from safetensors.torch import load_file
        m.load_state_dict(load_file(str(self.folder / "swar.safetensors")), strict=True)
        self.m = m.to(self.device).eval()
        self.mel = Mel().to(self.device)
        self._pause_ids = torch.tensor([STOI[c] for c in PAUSE_CHARS], device=self.device)
        self._pause_set = {STOI[c] for c in PAUSE_CHARS}
        self._voices: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}

    @classmethod
    def from_pretrained(cls, repo_or_path: str | Path = "elitexp/Swar", device: str | None = None,
                        revision: str | None = None) -> "Swar":
        p = Path(repo_or_path)
        if not p.is_dir():
            from huggingface_hub import snapshot_download
            p = Path(snapshot_download(str(repo_or_path), revision=revision))
        return cls(p, device)

    # ---- voices
    def list_voices(self) -> list[str]:
        return sorted(f.stem for f in (self.folder / "voices").glob("*.pt"))

    @torch.no_grad()
    def style_from_audio(self, path: str | Path) -> tuple[torch.Tensor, torch.Tensor]:
        import soundfile as sf
        import torchaudio.functional as AF
        wav, sr = sf.read(str(path), dtype="float32", always_2d=True)
        x = torch.from_numpy(wav.mean(1))
        if sr != SR:
            x = AF.resample(x, sr, SR)
        mel = self.mel(x.to(self.device)[None])[:, None]
        return self.m["style_encoder"](mel), self.m["predictor_encoder"](mel)

    def voice(self, voice: str | Path | None = None) -> tuple[torch.Tensor, torch.Tensor]:
        """A name from voices/, a .pt of style vectors, or an audio clip."""
        key = str(voice or "bijaya")
        if key not in self._voices:
            p = Path(key)
            if p.suffix.lower() in (".wav", ".flac", ".mp3", ".ogg"):
                self._voices[key] = self.style_from_audio(p)
            else:
                f = p if p.suffix == ".pt" else self.folder / "voices" / f"{key}.pt"
                v = torch.load(f, map_location="cpu", weights_only=True)
                self._voices[key] = (v["s_ac"].to(self.device), v["s_pr"].to(self.device))
        return self._voices[key]

    # ---- synthesis
    @torch.no_grad()
    def _tts(self, ids: list[int], s_ac, s_pr, speed: float, pause_floor: int,
             end_pause: int | None = None, tail: int = EDGE_UNITS) -> torch.Tensor:
        """end_pause: units for the sentence-final punctuation (None = as predicted); the
        last sentence gets a short one, so the clip does not end in half a second of nothing."""
        m, dev = self.m, self.device
        x_ids = torch.tensor([ids], device=dev)
        n = torch.tensor([len(ids)], device=dev)
        mask = torch.zeros(1, len(ids), dtype=torch.bool, device=dev)
        bert = m["plbert"](x_ids, attention_mask=(~mask).int()).last_hidden_state
        d_en = m["bert_encoder"](bert).transpose(1, 2)
        d = m["predictor"].text_encoder(d_en, s_pr, n.cpu(), mask)
        h, _ = m["predictor"].lstm(d)
        dur = torch.round(torch.sigmoid(m["predictor"].duration_proj(h)).sum(-1) / speed).clamp(min=1).long()
        if pause_floor:
            L = len(ids)
            pos = torch.arange(L, device=dev)
            p = torch.isin(x_ids[0], self._pause_ids) & (pos > 0) & (pos < L - 2)   # inside the sentence
            dur[0, p] = dur[0, p].clamp(min=pause_floor)
        dur[0, 0] = dur[0, 0].clamp(max=EDGE_UNITS)                 # room silence before the sentence
        dur[0, -1] = dur[0, -1].clamp(max=tail)                      # ... and after it
        if end_pause is not None and ids[-2] in self._pause_set:
            dur[0, -2] = max(1, end_pause)
        aln = _alignment(dur, int(dur.sum()))
        F0, N = m["predictor"].F0Ntrain(d.transpose(-1, -2) @ aln, s_pr)
        asr = m["text_encoder"](x_ids, n.cpu(), mask) @ aln
        return m["decoder"](asr, F0, N, s_ac).squeeze().clamp(-1, 1)

    def synthesize(self, text: str, voice: str | Path | None = None, speed: float = 1.0,
                   pause_floor_ms: int = 0, sentence_pause_ms: int = 350, loanwords: bool = True,
                   return_info: bool = False):
        """-> (float32 waveform, 16000), plus a timing dict if return_info."""
        t0 = time.time()
        s_ac, s_pr = self.voice(voice)
        floor = int(round(pause_floor_ms / 25))
        speed = float(min(2.0, max(0.5, speed)))
        gap = int(round(max(0, sentence_pause_ms) / 1000 / UNIT_S)) - 2 * EDGE_UNITS   # edges are silent too
        # split the RAW text first, so each sentence's numbers are read in its own
        # language ("Rs. 250" stays together: no split before a digit)
        raws = re.split(r"(?<=[।?!])\s+|(?<=\.)\s+(?![0-9०-९])", text)
        # loanwords: "recharge" inside a Nepali sentence is said रिचार्ज (swar/loanwords.py)
        sentences = [s for raw in raws for s in split_sentences(normalize(nepalise(raw) if loanwords else raw))]
        if not sentences:
            raise ValueError("nothing speakable left after normalising the text")
        pieces, first = [], None
        for k, sent in enumerate(sentences):
            last = k == len(sentences) - 1
            y = self._tts(to_ids(sent), s_ac, s_pr, speed, floor, end_pause=TAIL_UNITS if last else gap,
                          tail=TAIL_UNITS if last else EDGE_UNITS).float().cpu().numpy()
            first = first if first is not None else time.time() - t0
            pieces.append(y)
            if not last and sent[-1] not in PAUSE_CHARS:            # cut mid run-on: pause as silence
                pieces.append(np.zeros(int(max(0, gap) * UNIT_S * SR), np.float32))
        wav = np.concatenate(pieces)
        if not return_info:
            return wav, SR
        return wav, SR, {"first_audio_sec": round(first, 3), "total_sec": round(time.time() - t0, 3),
                         "duration_sec": round(len(wav) / SR, 2), "sentences": sentences}


def main() -> int:
    ap = argparse.ArgumentParser(description="Swar text-to-speech")
    ap.add_argument("text")
    ap.add_argument("-o", "--out", default="swar.wav")
    ap.add_argument("--model", default=str(Path(__file__).resolve().parents[1]))
    ap.add_argument("--voice", default=None, help="voice name, .pt, or a reference audio clip")
    ap.add_argument("--speed", type=float, default=1.0)
    ap.add_argument("--pause-floor-ms", type=int, default=0)
    ap.add_argument("--sentence-pause-ms", type=int, default=350)
    ap.add_argument("--device", default=None)
    a = ap.parse_args()
    import soundfile as sf
    tts = Swar.from_pretrained(a.model, device=a.device)
    wav, sr, info = tts.synthesize(a.text, voice=a.voice, speed=a.speed,
                                   pause_floor_ms=a.pause_floor_ms, sentence_pause_ms=a.sentence_pause_ms,
                                   return_info=True)
    sf.write(a.out, wav, sr)
    print(f"{a.out}: {info['duration_sec']} s audio in {info['total_sec']} s -- read as: "
          f"{' | '.join(info['sentences'])}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())