from __future__ import annotations

import logging
import re
from pathlib import Path
from typing import Optional

import numpy as np

from config import kokoro_bundle_dir, settings

logger = logging.getLogger(__name__)


def _kokoro_root() -> Path:
    p = kokoro_bundle_dir().resolve()
    if not p.is_dir():
        raise FileNotFoundError(f"Kokoro model directory missing: {p}")
    cfg = p / "config.json"
    w0 = p / "kokoro-v1_0.pth"
    w1 = p / "kokoro-v1_1-zh.pth"
    weight = w0 if w0.is_file() else (w1 if w1.is_file() else None)
    if not cfg.is_file() or weight is None:
        raise FileNotFoundError(f"Kokoro bundle incomplete (need config.json and weight .pth): {p}")
    return p


class TTSService:
    def __init__(self) -> None:
        self._pipeline = None
        self._kokoro_root: Optional[Path] = None
        self._ready = False
        self._attempted = False
        self._init_error: Optional[str] = None

    def _voice_pt_path(self) -> Path:
        root = self._kokoro_root
        if root is None:
            raise RuntimeError("Kokoro pipeline not initialized")
        name = (settings.tts_voice or "").strip()
        if not name:
            raise ValueError("TTS_VOICE must be set (voice name, e.g. zf_xiaoxiao)")
        p = root / "voices" / f"{name}.pt"
        if not p.is_file():
            raise FileNotFoundError(f"Voice file missing: {p}")
        return p

    def _ensure_loaded(self) -> None:
        if self._ready or self._attempted:
            return
        self._attempted = True
        try:
            import kokoro

            root = _kokoro_root()
            w0 = root / "kokoro-v1_0.pth"
            w1 = root / "kokoro-v1_1-zh.pth"
            weight = w0 if w0.is_file() else w1
            hf_repo_id = (
                "hexgrad/Kokoro-82M-v1.1-zh" if weight.name.startswith("kokoro-v1_1") else "hexgrad/Kokoro-82M"
            )
            kmodel = kokoro.KModel(
                repo_id=hf_repo_id,
                config=str(root / "config.json"),
                model=str(weight),
            )
            self._pipeline = kokoro.KPipeline(lang_code="z", model=kmodel, repo_id=str(root))
            self._kokoro_root = root
            self._ready = True
        except Exception as exc:  # pragma: no cover
            self._init_error = str(exc)
            logger.error("TTS init failed: %s", exc)

    @staticmethod
    def _normalize_text(text: str) -> str:
        cleaned = (text or "").strip()
        if not cleaned:
            return ""
        cleaned = cleaned.replace("\r", "\n")
        cleaned = re.sub(r"[\t\f\v]+", " ", cleaned)
        cleaned = re.sub(r"\s*([，。！？；：、,.!?;:])\s*", r"\1", cleaned)
        cleaned = re.sub(r"(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff])", "", cleaned)
        cleaned = re.sub(r"\n{2,}", "\n", cleaned)
        cleaned = re.sub(r"\s{2,}", " ", cleaned)
        return cleaned.strip()

    @staticmethod
    def _apply_edge_fade(audio: np.ndarray, sr: int) -> np.ndarray:
        if audio.size == 0:
            return audio
        fade_samples = max(1, min(int(sr * max(settings.tts_fade_ms, 0) / 1000.0), audio.shape[0] // 8))
        if fade_samples <= 1:
            return audio
        out = np.array(audio, copy=True)
        ramp = np.linspace(0.0, 1.0, fade_samples, dtype=np.float32)
        out[:fade_samples] *= ramp
        out[-fade_samples:] *= ramp[::-1]
        return out

    async def synthesize(self, text: str) -> tuple[np.ndarray, int]:
        text = self._normalize_text(text)
        if not text:
            return np.zeros(1, dtype=np.float32), 24000
        self._ensure_loaded()
        if not self._ready or self._pipeline is None:
            return np.zeros(1, dtype=np.float32), 24000
        try:
            voice_path = str(self._voice_pt_path())
        except (OSError, ValueError, RuntimeError) as exc:
            logger.error("TTS voice: %s", exc)
            return np.zeros(1, dtype=np.float32), 24000
        chunks = []
        for _, _, audio in self._pipeline(
            text,
            voice=voice_path,
            speed=max(0.8, min(1.1, settings.tts_speed)),
            split_pattern=r"\n+",
        ):
            chunks.append(self._apply_edge_fade(np.asarray(audio, dtype=np.float32), 24000))
        if chunks:
            return np.concatenate(chunks), 24000
        return np.zeros(1, dtype=np.float32), 24000

    @property
    def health(self) -> dict:
        return {"ready": self._ready, "attempted": self._attempted, "error": self._init_error}