"""MeloTTS 엔진 (in-process).

- 한국어/영어 지원, GPU 자동 사용, 속도/피치 조절
- 자동 모드(AUTO): 입력 텍스트를 언어별로 잘라 각 언어에 맞는 음성으로 합성 후 이어붙임
- 판별 불가(숫자/기호만 등) 구간의 기본 언어는 영어
"""
from __future__ import annotations

import io
import os
import re
import shutil
import subprocess
import tempfile
import threading

import numpy as np
import soundfile as sf
import torch

try:
    import librosa  # 피치 조절 / 폴백 타임스트레치
except Exception:  # pragma: no cover
    librosa = None

try:
    import pyrubberband as prb  # 고품질 타임스트레치(발음 보존)
except Exception:  # pragma: no cover
    prb = None

from melo.api import TTS

from .base import BaseEngine
from .ko_normalize import normalize_korean

# 영어 화자(억양) 목록
EN_SPEAKERS = [
    {"id": "EN-US", "label": "영어 · 미국식"},
    {"id": "EN-BR", "label": "영어 · 영국식"},
    {"id": "EN_INDIA", "label": "영어 · 인도식"},
    {"id": "EN-AU", "label": "영어 · 호주식"},
    {"id": "EN-Default", "label": "영어 · 기본"},
]
KR_SPEAKERS = [{"id": "KR", "label": "한국어 기본 목소리"}]
JP_SPEAKERS = [{"id": "JP", "label": "일본어 기본 목소리"}]
ZH_SPEAKERS = [{"id": "ZH", "label": "중국어 기본 목소리"}]

# 실제 MeloTTS 언어 코드 (한/영/일/중)
MELO_LANG = {"KR": "KR", "EN": "EN", "JP": "JP", "ZH": "ZH"}

DEFAULT_LANG = "EN"  # 판별 불가 구간의 기본 언어

# 문장 경계: 종결부호 뒤 공백, 또는 줄바꿈
_SENT_SPLIT_RE = re.compile(r"(?<=[.!?。！？…])\s+|\n+")


def split_sentences(text: str) -> list[str]:
    """텍스트를 문장 단위로 분리(종결부호/줄바꿈 기준)."""
    parts = [p.strip() for p in _SENT_SPLIT_RE.split(text) if p and p.strip()]
    return parts or ([text.strip()] if text.strip() else [])

_KR_RE = re.compile(r"[가-힣ᄀ-ᇿ㄰-㆏]")
_EN_RE = re.compile(r"[A-Za-z]")
# 일본어 가나(히라가나/가타카나/반각 가타카나)
_KANA_RE = re.compile(r"[぀-ヿｦ-ﾝ]")
# 한자(CJK 통합 한자) — 일본어/중국어 공유라 단독으로는 모호
_HAN_RE = re.compile(r"[㐀-䶿一-鿿豈-﫿]")


def _pick_device() -> str:
    return "cuda:0" if torch.cuda.is_available() else "cpu"


def _char_group(ch: str) -> str | None:
    """문자를 그룹으로 분류: KR / EN / CJK(가나+한자) / None(중립)."""
    if _KR_RE.match(ch):
        return "KR"
    if _KANA_RE.match(ch) or _HAN_RE.match(ch):
        return "CJK"
    if _EN_RE.match(ch):
        return "EN"
    return None


def split_by_language(text: str, default_lang: str = DEFAULT_LANG) -> list[tuple[str, str]]:
    """입력을 언어별(한/영/일/중) 구간으로 분리.

    - 한글→KR, 라틴문자→EN
    - 가나+한자는 하나의 CJK 구간으로 묶은 뒤, 가나가 있으면 일본어(JP), 없으면 중국어(ZH)로 판정
    - 숫자·기호·공백 등 중립 문자는 진행 중 구간에 붙임
    - 전체가 판별 불가면 default_lang(기본 영어)
    """
    raw: list[tuple[str, str]] = []
    cur_group: str | None = None
    cur: list[str] = []

    for ch in text:
        g = _char_group(ch)
        if g is None:
            cur.append(ch)
            continue
        if cur_group is None:
            cur_group = g
            cur.append(ch)
        elif g == cur_group:
            cur.append(ch)
        else:
            raw.append((cur_group, "".join(cur)))
            cur_group = g
            cur = [ch]

    if cur:
        raw.append((cur_group if cur_group is not None else "DEFAULT", "".join(cur)))

    out: list[tuple[str, str]] = []
    for grp, seg in raw:
        if not seg.strip():
            continue
        if grp == "CJK":
            lang = "JP" if _KANA_RE.search(seg) else "ZH"
        elif grp in ("KR", "EN"):
            lang = grp
        else:
            lang = default_lang
        out.append((lang, seg))
    return out


class MeloEngine(BaseEngine):
    id = "melo"
    label = "MeloTTS"
    license = "MIT"
    uses_gpu = True
    notes = "자동 언어 감지: 한국어/영어/일본어/중국어 혼합 텍스트를 언어별로 나눠 각 음성으로 합성. 기본 영어."

    def __init__(self) -> None:
        self.device = _pick_device()
        self._models: dict[str, TTS] = {}
        self._locks: dict[str, threading.Lock] = {k: threading.Lock() for k in MELO_LANG}
        self._global_lock = threading.Lock()

    # ---- 모델 로딩 -------------------------------------------------
    def _get_model(self, lang: str) -> TTS:
        if lang not in MELO_LANG:
            raise ValueError(f"지원하지 않는 언어입니다: {lang}")
        model = self._models.get(lang)
        if model is not None:
            return model
        with self._locks[lang]:
            model = self._models.get(lang)
            if model is None:
                model = TTS(language=MELO_LANG[lang], device=self.device)
                self._models[lang] = model
            return model

    def warmup(self) -> None:
        for lang in MELO_LANG:
            try:
                self._get_model(lang)
            except Exception as exc:  # pragma: no cover
                print(f"[melo warmup] {lang} 모델 로드 실패: {exc}")

    def describe(self) -> dict:
        languages = [
            {
                "code": "AUTO",
                "label": "자동 감지 (한/영/일/중)",
                "speakers": EN_SPEAKERS,  # 영어 구간에 쓸 억양 선택 (그 외 언어는 자동)
            },
            {"code": "EN", "label": "English (영어)", "speakers": EN_SPEAKERS},
            {"code": "KR", "label": "한국어", "speakers": KR_SPEAKERS},
            {"code": "JP", "label": "日本語 (일본어)", "speakers": JP_SPEAKERS},
            {"code": "ZH", "label": "中文 (중국어)", "speakers": ZH_SPEAKERS},
        ]
        return {
            "id": self.id,
            "label": self.label,
            "license": self.license,
            "uses_gpu": self.uses_gpu,
            "notes": self.notes,
            "languages": languages,
            "default_language": "AUTO",
            "supports": {
                "speed": True,
                "pitch": True,
                "word_gap": True,
                "sentence_gap": True,
            },
        }

    def _resolve_speaker(self, model: TTS, lang: str, speaker: str | None) -> int:
        spk2id = model.hps.data.spk2id
        if speaker and speaker in spk2id:
            return spk2id[speaker]
        default_by_lang = {"KR": "KR", "EN": "EN-US", "JP": "JP", "ZH": "ZH"}
        default_id = default_by_lang.get(lang)
        if default_id and default_id in spk2id:
            return spk2id[default_id]
        return list(spk2id.values())[0]

    def _synth_segment(
        self, text: str, lang: str, speaker: str | None
    ) -> tuple[np.ndarray, int]:
        # 항상 자연 속도(1.0)로 합성 — 속도는 사후 타임스트레치로 처리해 발음 유지
        model = self._get_model(lang)
        speaker_id = self._resolve_speaker(model, lang, speaker)
        with self._global_lock:
            audio = model.tts_to_file(
                text, speaker_id, output_path=None, speed=1.0, quiet=True
            )
        sr = model.hps.data.sampling_rate
        return np.asarray(audio, dtype=np.float32), sr

    def _time_stretch(self, audio: np.ndarray, sr: int, rate: float) -> np.ndarray:
        """오디오를 rate 배속으로 변경(피치 보존). rate>1 = 빠르게.

        자연 속도(1.0)로 합성한 뒤 여기서 타임스트레치를 적용해 글자 발화
        속도만 조절한다(피치 보존).
        우선순위: ffmpeg atempo(기본) → Rubberband(R3) → librosa(phase vocoder).
        """
        if abs(rate - 1.0) < 1e-3:
            return audio
        audio = np.asarray(audio, dtype=np.float32)

        # 1) ffmpeg atempo (기본값: 추가 의존성 없음, 무난)
        if shutil.which("ffmpeg") and 0.5 <= rate <= 2.0:
            try:
                return self._atempo(audio, sr, rate)
            except Exception:
                pass

        # 2) Rubberband R3 (발음/포먼트 보존 폴백)
        if prb is not None and shutil.which("rubberband"):
            try:
                out = prb.time_stretch(
                    audio, sr, rate, rbargs={"--engine": "finer"}
                )
                return np.asarray(out, dtype=np.float32)
            except Exception:
                try:
                    out = prb.time_stretch(audio, sr, rate)
                    return np.asarray(out, dtype=np.float32)
                except Exception:
                    pass

        # 3) librosa phase vocoder (최후 폴백)
        if librosa is not None:
            try:
                return librosa.effects.time_stretch(audio, rate).astype(np.float32)
            except Exception:
                try:
                    return librosa.effects.time_stretch(
                        audio, rate=rate
                    ).astype(np.float32)
                except Exception:
                    pass
        return audio

    @staticmethod
    def _atempo(audio: np.ndarray, sr: int, rate: float) -> np.ndarray:
        with tempfile.TemporaryDirectory() as d:
            src = os.path.join(d, "in.wav")
            dst = os.path.join(d, "out.wav")
            sf.write(src, audio, sr, format="WAV", subtype="PCM_16")
            subprocess.run(
                [
                    "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
                    "-i", src, "-filter:a", f"atempo={rate:.4f}", dst,
                ],
                check=True,
            )
            out, _ = sf.read(dst, dtype="float32")
            return np.asarray(out, dtype=np.float32)

    def _synth_natural(
        self, text: str, language: str, speaker: str | None
    ) -> tuple[np.ndarray, int]:
        """자연 속도(1.0)로 텍스트 한 조각을 합성. AUTO 모드는 언어별로 나눠 합성.

        속도/문장·단어 간격 조절은 이 결과 위에서 처리한다(여기서는 하지 않음).
        """
        if language == "AUTO":
            # 영어 구간엔 사용자가 고른 억양, 한국어 구간엔 KR 음성
            en_speaker = speaker if (speaker and speaker.startswith("EN")) else "EN-US"
            spk_by_lang = {"KR": "KR", "EN": en_speaker, "JP": "JP", "ZH": "ZH"}
            parts = split_by_language(text, default_lang=DEFAULT_LANG)
            if not parts:
                parts = [(DEFAULT_LANG, text)]
            audios: list[np.ndarray] = []
            sr = None
            gap = None
            for seg_lang, seg_text in parts:
                seg_spk = spk_by_lang.get(seg_lang, en_speaker)
                audio, seg_sr = self._synth_segment(seg_text, seg_lang, seg_spk)
                if sr is None:
                    sr = seg_sr
                    gap = np.zeros(int(sr * 0.06), dtype=np.float32)
                elif seg_sr != sr and librosa is not None:
                    audio = librosa.resample(
                        audio, orig_sr=seg_sr, target_sr=sr
                    ).astype(np.float32)
                if audios:
                    audios.append(gap)
                audios.append(audio)
            audio = np.concatenate(audios) if audios else np.zeros(1, dtype=np.float32)
            return audio, (sr or self._get_model(DEFAULT_LANG).hps.data.sampling_rate)
        return self._synth_segment(text, language, speaker)

    # ---- 합성 -----------------------------------------------------
    def synth_wav(
        self,
        text: str,
        language: str,
        speaker: str | None = None,
        speed: float = 1.0,
        pitch: float = 0.0,
        word_gap: float = 0.0,
        sentence_gap: float | None = None,
    ) -> bytes:
        """세 가지를 독립적으로 조절:
        - speed: 글자(음소) 발화 속도 — 자연 합성 후 타임스트레치(피치 보존)
        - word_gap: 단어 사이 추가 무음(초). >0 이면 어절 단위로 나눠 합성
        - sentence_gap: 문장 사이 무음(초)
        간격은 타임스트레치의 영향을 받지 않아 글자속도와 독립적으로 유지된다.
        """
        text = (text or "").strip()
        if not text:
            raise ValueError("텍스트가 비어 있습니다.")

        speed = float(max(0.5, min(2.0, speed)))
        pitch = float(max(-12.0, min(12.0, pitch)))
        word_gap = float(max(0.0, min(1.0, word_gap)))
        sentence_gap = 0.30 if sentence_gap is None else sentence_gap
        sentence_gap = float(max(0.0, min(2.0, sentence_gap)))
        language = (language or "AUTO").upper()

        # 한국어 발음 개선: 기호/약어/사용자사전 전처리 (한국어·자동 모드)
        if language in ("AUTO", "KR"):
            text = normalize_korean(text)

        def stretch(a: np.ndarray, sr: int) -> np.ndarray:
            if abs(speed - 1.0) > 1e-3:
                return self._time_stretch(a, sr, speed)
            return a

        sentences = split_sentences(text) or [text]
        sr: int | None = None
        pieces: list[np.ndarray] = []

        for si, sent in enumerate(sentences):
            if word_gap > 1e-4:
                # 어절 단위로 나눠 합성 후 단어 간 무음 삽입
                words = sent.split()
                wpieces: list[np.ndarray] = []
                for wi, word in enumerate(words):
                    a, s = self._synth_natural(word, language, speaker)
                    if sr is None:
                        sr = s
                    a = stretch(a, sr)
                    if wi > 0:
                        wpieces.append(np.zeros(int(sr * word_gap), dtype=np.float32))
                    wpieces.append(a)
                sent_audio = (
                    np.concatenate(wpieces)
                    if wpieces
                    else np.zeros(1, dtype=np.float32)
                )
            else:
                a, s = self._synth_natural(sent, language, speaker)
                if sr is None:
                    sr = s
                sent_audio = stretch(a, sr)

            if pieces and sentence_gap > 1e-4:
                pieces.append(np.zeros(int(sr * sentence_gap), dtype=np.float32))
            pieces.append(sent_audio)

        audio = np.concatenate(pieces) if pieces else np.zeros(1, dtype=np.float32)
        if sr is None:
            sr = self._get_model(DEFAULT_LANG).hps.data.sampling_rate

        if abs(pitch) > 1e-3 and librosa is not None:
            audio = librosa.effects.pitch_shift(audio, sr, n_steps=pitch)

        buf = io.BytesIO()
        sf.write(buf, audio, sr, format="WAV", subtype="PCM_16")
        buf.seek(0)
        return buf.read()