app/core/cloud_tts.py

"""
app/core/cloud_tts.py
──────────────────────
High-Performance Parallel Cloud TTS Module using Microsoft Edge-TTS.
Features:
1. Concurrent synthesis with asyncio.gather + Semaphore(8).
2. Vietnamese text normalization (numbers, dates, currency, abbreviations).
3. Dynamic Audio Time-Stretching (FFmpeg atempo): Auto-speeds up sentences if TTS duration exceeds timeline slot (up to 1.65x).
4. Full Audio Integrity: NEVER hard-slices or cuts off ending syllables/words.
5. Exact millisecond timeline synchronization.
"""

import os
import re
import json
import time
import asyncio
import subprocess
from pathlib import Path
from typing import List, Dict, Optional, Callable

from app.core.vietnamese_text_normalizer import VietnameseTextNormalizer

# ── OmniVoice presets — sync với D:\omnivoice - web\app.py ──
OMNIVOICE_PRESETS = {
    "fun_male": "male, young adult, high pitch",
    "story_male": "male, young adult, moderate pitch",
    "meme_male": "male, young adult, moderate pitch",
    "calm_fun_male": "male, young adult, low pitch",
}
OMNIVOICE_EXTENDED = {"meme_male_gasp", "meme_male_normal", "meme_male_excited", "meme_male_confident",
                       "fun_male_gasp", "story_male_normal", "calm_fun_male_sad", "calm_fun_male_whispering",
                       "meme_male_sad", "meme_male_whispering", "fun_male_excited"}
OMNIVOICE_ALL = set(OMNIVOICE_PRESETS.keys()) | OMNIVOICE_EXTENDED

def _is_omnivoice_voice(v: str) -> bool:
    v = str(v or "").strip()
    if v in OMNIVOICE_ALL:
        return True
    for p in OMNIVOICE_PRESETS:
        if v.startswith(p):
            return True
    return False

def _resolve_omnivoice_preset_cloud(voice: str) -> str:
    v = str(voice or "").strip()
    if v in OMNIVOICE_PRESETS:
        return v
    for p in OMNIVOICE_PRESETS:
        if v.startswith(p):
            return p
    return "meme_male"

# Placeholder dịch fail — TTS phải SKIP chứ không đọc lên thành tiếng
TTS_PLACEHOLDER_NORMED = {"CẦN DỊCH LẠI", "CAN DICH LAI", "DỊCH LẠI", "CẦN DỊCH"}

def _is_tts_placeholder(text: str) -> bool:
    import re as _re
    t = str(text or "").strip()
    if not t:
        return False
    norm = _re.sub(r"\s+", " ", t).strip().strip("[]").upper()
    return norm in TTS_PLACEHOLDER_NORMED


class CloudTTSEngine:
    def __init__(self, log_fn: Optional[Callable[[str], None]] = None, ffmpeg_path: str = "ffmpeg"):
        self.log_fn = log_fn or print
        self.ffmpeg_path = ffmpeg_path
        self.normalizer = VietnameseTextNormalizer()
        # OmniVoice Cloud config — đọc từ config.json hoặc env
        self.omnivoice_api_url = os.getenv("OMNIVOICE_API_URL", "https://hoangtaiii-omnivoice.hf.space").rstrip("/")
        self.omnivoice_api_key = os.getenv("OMNIVOICE_API_KEY", "sk-demo123")
        # HF token cho ZeroGPU quota (https://huggingface.co/settings/tokens) — ưu tiên env HF_TOKEN
        self.hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN") or os.getenv("HUGGINGFACE_TOKEN") or ""
        # Thử đọc config.json nếu có
        try:
            cfg_path = Path(__file__).resolve().parents[2] / "config.json"
            if cfg_path.exists():
                cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
                self.hf_token = self.hf_token or cfg.get("hf_key", "") or cfg.get("hf_token", "")
                tts_cfg = cfg.get("tts", {})
                self.omnivoice_api_url = (tts_cfg.get("omnivoice_api_url") or tts_cfg.get("omnivoice_cloud_url") or self.omnivoice_api_url).rstrip("/")
                self.omnivoice_api_key = tts_cfg.get("omnivoice_api_key") or tts_cfg.get("omnivoice_cloud_api_key") or self.omnivoice_api_key
                # ưu tiên token trong tts config nếu có
                self.hf_token = self.hf_token or tts_cfg.get("hf_token", "") or tts_cfg.get("hf_key", "")
        except Exception:
            pass

    def _log(self, msg: str):
        self.log_fn(f"[Cloud TTS] {msg}")

    def synthesize_srt_to_audio(
        self,
        srt_path: str,
        output_wav: str,
        voice: str = "vi-VN-NamMinhNeural",
        speed: float = 1.0,
        pitch: int = 0,
        volume: int = 100,
        temp_dir: Optional[str] = None
    ) -> bool:
        srt_file = Path(srt_path)
        if not srt_file.exists():
            self._log(f"❌ SRT file not found: {srt_path}")
            return False

        out_file = Path(output_wav)
        out_file.parent.mkdir(parents=True, exist_ok=True)
        
        segments_dir = Path(temp_dir or (out_file.parent / "tts_segments"))
        segments_dir.mkdir(parents=True, exist_ok=True)

        blocks = self._parse_srt(srt_file)
        if not blocks:
            self._log("⚠️ Không có block phụ đề nào để tổng hợp TTS.")
            return False

        # ── Auto-dispatch: OmniVoice Cloud nếu voice là preset mới ──
        is_ov = _is_omnivoice_voice(voice)
        if is_ov:
            self._log(f"🎙️ Phát hiện giọng OmniVoice Cloud: {voice} -> {self.omnivoice_api_url}")
            try:
                import requests  # check sẵn
            except ImportError:
                self._log("❌ requests chưa cài — pip install requests")
                return False
            self._log(f"⚡ Tạo giọng đọc OmniVoice Cloud song song cho {len(blocks)} câu ({voice})...")
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                success = loop.run_until_complete(
                    self._synthesize_blocks_parallel_omnivoice(blocks, segments_dir, voice, speed)
                )
            finally:
                loop.close()
        else:
            self._log(f"⚡ Tạo giọng đọc song song cho {len(blocks)} câu thoại ({voice})...")
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                success = loop.run_until_complete(
                    self._synthesize_blocks_parallel(blocks, segments_dir, voice, speed, pitch, volume)
                )
            finally:
                loop.close()

        if not success:
            self._log("❌ Lỗi tổng hợp giọng đọc TTS.")
            return False

        self._log("🎚️ Đang tự động co dãn tốc độ (Time-Stretch) & giữ trọn vẹn 100% âm thanh cuối câu...")
        return self._merge_segments_to_timeline(blocks, segments_dir, out_file)

    async def _synthesize_blocks_parallel(
        self,
        blocks: List[Dict],
        segments_dir: Path,
        voice: str,
        speed: float,
        pitch: int,
        volume: int
    ) -> bool:
        import edge_tts

        rate_str = f"{int(round((speed - 1.0) * 100)):+d}%"
        pitch_str = f"{int(round(pitch)):+d}Hz"
        vol_str = f"{int(round(volume - 100)):+d}%"

        semaphore = asyncio.Semaphore(8)

        async def _synthesize_one(block: Dict):
            raw_text = block["text"].strip()
            if not raw_text:
                return
            # FIX: placeholder "[CẦN DỊCH LẠI]" mà đọc lên là bug P0 — skip im lặng
            if _is_tts_placeholder(raw_text):
                self._log(f"⏭️ Skip block {block['id']}: placeholder dịch fail, để im lặng.")
                return

            # Apply Vietnamese normalizer
            text = self.normalizer.normalize(raw_text) if hasattr(self.normalizer, "normalize") else raw_text
            text = self._sanitize_for_tts(text)

            seg_path = segments_dir / f"seg_{block['id']:04d}.mp3"
            wav_path = segments_dir / f"seg_{block['id']:04d}.wav"
            sidecar_path = segments_dir / f"seg_{block['id']:04d}.txt"

            # FIX stale cache: trước đây chỉ check wav tồn tại -> sửa sub không có tác dụng.
            # Nay so text sidecar, khác text là tổng hợp lại.
            if wav_path.exists() and wav_path.stat().st_size > 500:
                try:
                    if sidecar_path.exists() and sidecar_path.read_text(encoding="utf-8") == text:
                        return
                except Exception:
                    pass

            async with semaphore:
                for attempt in range(3):
                    try:
                        communicate = edge_tts.Communicate(
                            text,
                            voice=voice,
                            rate=rate_str,
                            pitch=pitch_str,
                            volume=vol_str
                        )
                        await communicate.save(str(seg_path))

                        # Convert to standard 16kHz mono WAV with gentle silence padding at end to prevent clipping
                        cmd = [
                            str(self.ffmpeg_path), "-y",
                            "-i", str(seg_path),
                            "-af", "apad=pad_dur=0.08",
                            "-ar", "16000", "-ac", "1",
                            str(wav_path)
                        ]
                        subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
                        if seg_path.exists():
                            seg_path.unlink()
                        try:
                            sidecar_path.write_text(text, encoding="utf-8")
                        except Exception:
                            pass
                        break
                    except Exception as e:
                        if attempt == 2:
                            self._log(f"⚠️ Thất bại block {block['id']}: {e}")
                        await asyncio.sleep(0.5)

        tasks = [_synthesize_one(b) for b in blocks]
        await asyncio.gather(*tasks)
        return True

    async def _synthesize_blocks_parallel_omnivoice(
        self,
        blocks: List[Dict],
        segments_dir: Path,
        voice: str,
        speed: float,
    ) -> bool:
        """FIX P0: hàm này được gọi ở synthesize_srt_to_audio nhưng chưa từng tồn tại
        -> mọi giọng OmniVoice trên cloud path rớt AttributeError -> job FAILED.
        Dùng FastAPI /v1/generate (thread-safe qua asyncio.to_thread), semaphore 4."""
        import requests as _requests

        try:
            speed_val = max(0.5, min(1.5, float(speed)))
        except Exception:
            speed_val = 1.08
        preset = _resolve_omnivoice_preset_cloud(voice)
        headers = {"X-API-Key": self.omnivoice_api_key, "Content-Type": "application/json"}
        semaphore = asyncio.Semaphore(4)

        def _post_one(payload: dict) -> bytes:
            resp = _requests.post(
                f"{self.omnivoice_api_url}/v1/generate",
                headers=headers, json=payload, timeout=240)
            if resp.status_code != 200:
                raise RuntimeError(f"OmniVoice Cloud {resp.status_code}: {resp.text[:300]}")
            ctype = resp.headers.get("content-type", "")
            if "audio" not in ctype and len(resp.content) < 1000:
                raise RuntimeError(f"OmniVoice Cloud non-audio: {resp.content[:300]!r}")
            return resp.content

        async def _synthesize_one_ov(block: Dict):
            raw_text = str(block.get("text", "") or "").strip()
            if not raw_text or _is_tts_placeholder(raw_text):
                return
            text = self.normalizer.normalize(raw_text) if hasattr(self.normalizer, "normalize") else raw_text
            text = self._sanitize_for_tts(text)
            if not text:
                return
            wav_path = segments_dir / f"seg_{block['id']:04d}.wav"
            sidecar_path = segments_dir / f"seg_{block['id']:04d}.txt"
            if wav_path.exists() and wav_path.stat().st_size > 500:
                try:
                    if sidecar_path.exists() and sidecar_path.read_text(encoding="utf-8") == text:
                        return
                except Exception:
                    pass
            payload = {"text": text, "preset": preset, "speed": float(speed_val),
                       "steps": 16, "guidance": 2.0, "language": "vi"}
            async with semaphore:
                for attempt in range(3):
                    try:
                        audio_bytes = await asyncio.to_thread(_post_one, payload)
                        tmp_raw = segments_dir / f"seg_{block['id']:04d}.raw"
                        tmp_raw.write_bytes(audio_bytes)
                        cmd = [str(self.ffmpeg_path), "-y", "-i", str(tmp_raw),
                               "-af", "apad=pad_dur=0.08",
                               "-ar", "16000", "-ac", "1", str(wav_path)]
                        subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
                        try:
                            tmp_raw.unlink()
                        except Exception:
                            pass
                        try:
                            sidecar_path.write_text(text, encoding="utf-8")
                        except Exception:
                            pass
                        break
                    except Exception as e:
                        if attempt == 2:
                            self._log(f"⚠️ OmniVoice block {block['id']} thất bại: {e}")
                        await asyncio.sleep(1.0)

        await asyncio.gather(*[_synthesize_one_ov(b) for b in blocks])
        return True

    def _merge_segments_to_timeline(self, blocks: List[Dict], segments_dir: Path, output_wav: Path) -> bool:
        try:
            from pydub import AudioSegment
        except ImportError:
            self._log("❌ pydub not installed.")
            return False

        if not blocks:
            return False

        sorted_blocks = sorted(blocks, key=lambda x: x["start_ms"])
        total_duration_ms = max(b["end_ms"] for b in sorted_blocks) + 10000
        canvas = AudioSegment.silent(duration=total_duration_ms, frame_rate=16000)

        for i, b in enumerate(sorted_blocks):
            wav_path = segments_dir / f"seg_{b['id']:04d}.wav"
            if not wav_path.exists() or wav_path.stat().st_size < 300:
                continue

            try:
                seg_audio = AudioSegment.from_wav(str(wav_path))
                actual_duration_ms = len(seg_audio)

                target_duration_ms = max(300, b["end_ms"] - b["start_ms"])
                if i + 1 < len(sorted_blocks):
                    next_start_ms = sorted_blocks[i + 1]["start_ms"]
                    gap_to_next = next_start_ms - b["start_ms"]
                    if gap_to_next > 200:
                        max_allowed_ms = gap_to_next
                    else:
                        max_allowed_ms = target_duration_ms
                else:
                    max_allowed_ms = target_duration_ms + 2500

                safety_margin_ms = 80  # Khoảng đệm an toàn giữa 2 câu
                max_allowed_ms = max(150, max_allowed_ms - safety_margin_ms)

                # Dynamic Time-Stretching (atempo) nếu TTS câu thoại dài hơn khung thời gian
                if actual_duration_ms > max_allowed_ms:
                    speed_ratio = actual_duration_ms / float(max_allowed_ms)
                    effective_ratio = min(speed_ratio, 2.5)

                    atempo_filters = []
                    rem = effective_ratio
                    while rem > 2.0:
                        atempo_filters.append("atempo=2.0")
                        rem /= 2.0
                    if rem > 1.01:
                        atempo_filters.append(f"atempo={rem:.3f}")

                    if atempo_filters:
                        filter_chain = ",".join(atempo_filters)
                        stretched_wav = segments_dir / f"seg_{b['id']:04d}_stretched.wav"
                        cmd_stretch = [
                            str(self.ffmpeg_path), "-y",
                            "-i", str(wav_path),
                            "-filter:a", filter_chain,
                            "-ar", "16000", "-ac", "1",
                            str(stretched_wav)
                        ]
                        res = subprocess.run(cmd_stretch, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                        if res.returncode == 0 and stretched_wav.exists() and stretched_wav.stat().st_size > 300:
                            seg_audio = AudioSegment.from_wav(str(stretched_wav))

                # FIX: KHÔNG cắt cứng đuôi (trái docstring "NEVER hard-slices").
                # Trước đây câu dài hơn slot bị xén chữ im lặng. Nay giữ trọn audio,
                # chỉ log cảnh báo để QA loop / người dùng rút gọn text.
                if len(seg_audio) > max_allowed_ms:
                    over_ms = len(seg_audio) - int(max_allowed_ms)
                    self._log(f"⚠️ Block {b['id']} TTS tràn slot {over_ms}ms sau stretch "
                              f"({len(seg_audio)}>{int(max_allowed_ms)}) — giữ trọn, có thể lấn nhẹ câu kế.")

                fade_ms = min(50, len(seg_audio))
                seg_audio = seg_audio.fade_out(fade_ms)

                # Overlay segment lên timeline đúng vị trí start_ms
                canvas = canvas.overlay(seg_audio, position=b["start_ms"])

            except Exception as e:
                self._log(f"⚠️ Error overlaying block {b['id']}: {e}")

        canvas.export(str(output_wav), format="wav")
        return True

    def _sanitize_for_tts(self, text: str) -> str:
        replacements = {
            "\u2011": "-",
            "\u2010": "-",
            "\u2013": "-",
            "\u2014": "-",
            "\u200b": "",
            "\ufeff": "",
        }
        for src, dst in replacements.items():
            text = text.replace(src, dst)
        return text.strip()

    def _parse_srt(self, srt_path: Path) -> List[Dict]:
        content = srt_path.read_text(encoding="utf-8", errors="ignore")
        # FIX: giờ cho phép 1 chữ số (\d{1,2}) — Whisper/tool lẻ ra "0:00:01,000" bị miss cả block
        pattern = r"(\d+)\s+(\d{1,2}:\d{2}:\d{2}[.,]\d{3})\s*-->\s*(\d{1,2}:\d{2}:\d{2}[.,]\d{3})\s*\n(.*?)(?=\n\s*\d+\s+\d{1,2}:\d{2}:\d{2}[.,]\d{3}\s*-->|\Z)"
        
        blocks = []
        for m in re.finditer(pattern, content, re.DOTALL):
            b_id = int(m.group(1))
            start_str = m.group(2).strip().replace(".", ",")
            end_str = m.group(3).strip().replace(".", ",")
            text = " ".join(line.strip() for line in m.group(4).splitlines() if line.strip())
            
            if text:
                blocks.append({
                    "id": b_id,
                    "start_str": start_str,
                    "end_str": end_str,
                    "start_ms": self._ts_to_ms(start_str),
                    "end_ms": self._ts_to_ms(end_str),
                    "text": text
                })
        return blocks

    def _ts_to_ms(self, ts: str) -> int:
        ts = ts.strip().replace(".", ",")
        m = re.match(r"(\d+):(\d+):(\d+)[,](\d+)", ts)
        if m:
            h, mins, s, ms = map(int, m.groups())
            return ((h * 3600 + mins * 60 + s) * 1000) + ms
        return 0