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

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

import io
import re
import threading

import numpy as np
import soundfile as sf
import torch

try:
    import librosa  # 피치 조절용
except Exception:  # pragma: no cover
    librosa = None

from melo.api import TTS

from .base import BaseEngine

# 영어 화자(억양) 목록
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": "한국어 기본 목소리"}]

# 실제 MeloTTS 언어 코드
MELO_LANG = {"KR": "KR", "EN": "EN"}

DEFAULT_LANG = "EN"  # 판별 불가 구간의 기본 언어(사용자 요청: 기본값 영어)

_KR_RE = re.compile(r"[가-힣ᄀ-ᇿ㄰-㆏]")
_EN_RE = re.compile(r"[A-Za-z]")


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


def split_by_language(text: str, default_lang: str = DEFAULT_LANG) -> list[tuple[str, str]]:
    """텍스트를 한국어/영어 구간으로 분리.

    숫자·기호·공백 등 언어 판별이 안 되는 문자는 인접(진행 중) 구간에 붙인다.
    전체가 판별 불가면 default_lang 으로 처리한다.
    """
    segments: list[tuple[str, str]] = []
    cur_lang: str | None = None
    cur: list[str] = []

    for ch in text:
        if _KR_RE.match(ch):
            kind = "KR"
        elif _EN_RE.match(ch):
            kind = "EN"
        else:
            kind = None  # 중립 문자

        if kind is None:
            cur.append(ch)
            continue
        if cur_lang is None:
            cur_lang = kind
            cur.append(ch)
        elif kind == cur_lang:
            cur.append(ch)
        else:
            segments.append((cur_lang, "".join(cur)))
            cur_lang = kind
            cur = [ch]

    if cur:
        lang = cur_lang if cur_lang is not None else default_lang
        segments.append((lang, "".join(cur)))

    # 실제 발음 가능한(문자 포함) 구간만 유지
    return [(l, t) for l, t in segments if t.strip()]


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},
        ]
        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},
        }

    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_id = "KR" if lang == "KR" else "EN-US"
        if default_id in spk2id:
            return spk2id[default_id]
        return list(spk2id.values())[0]

    def _synth_segment(
        self, text: str, lang: str, speaker: str | None, speed: float
    ) -> tuple[np.ndarray, int]:
        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=speed, quiet=True
            )
        sr = model.hps.data.sampling_rate
        return np.asarray(audio, dtype=np.float32), sr

    # ---- 합성 -----------------------------------------------------
    def synth_wav(
        self,
        text: str,
        language: str,
        speaker: str | None = None,
        speed: float = 1.0,
        pitch: float = 0.0,
    ) -> bytes:
        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)))
        language = (language or "AUTO").upper()

        if language == "AUTO":
            # 영어 구간엔 사용자가 고른 억양, 한국어 구간엔 KR 음성
            en_speaker = speaker if (speaker and speaker.startswith("EN")) else "EN-US"
            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 = "KR" if seg_lang == "KR" else en_speaker
                audio, seg_sr = self._synth_segment(seg_text, seg_lang, seg_spk, speed)
                if sr is None:
                    sr = seg_sr
                    gap = np.zeros(int(sr * 0.06), dtype=np.float32)
                elif seg_sr != sr:
                    if 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)
        else:
            audio, sr = self._synth_segment(text, language, speaker, speed)

        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()