#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""FIKO AI - Yerel TTS sunucusu (Chatterbox Multilingual, Turkce).

OpenAI uyumlu uc: POST /v1/audio/speech
  Istek : JSON {model, input, voice, response_format, speed}
  Yanit : ham WAV baytlari (response_format ne olursa olsun WAV doner;
          backend icin .env'de TTS_FORMAT=wav olmali)

Ses secimi (voice alani):
  - models/voices/<voice>.wav varsa o referans klip kullanilir
  - yoksa models/voices/referans.wav (varsayilan kadin sesi klibi)
  - o da yoksa modelin yerlesik varsayilan sesi

Duygu: --exaggeration (0..1). Duygunun buyuk kismi referans klibin tonundan gelir;
canli ve ifadeli okunmus bir klip daha duygulu sonuc verir.

Kullanim:
  python tools/tts_server.py 5000
"""
import argparse
import io
import os
import re
import sys
import threading
import time
import wave

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
HF_HOME = os.path.join(ROOT, "models", "hf-cache")
os.environ.setdefault("HF_HOME", HF_HOME)
os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")
# Model onbellekte varsa agi hic kullanma (huggingface_hub bu ayari import aninda okur).
if os.path.isdir(os.path.join(HF_HOME, "hub", "models--ResembleAI--chatterbox", "snapshots")):
    os.environ.setdefault("HF_HUB_OFFLINE", "1")

import numpy as np  # noqa: E402
import torch  # noqa: E402
from flask import Flask, Response, jsonify, request  # noqa: E402
from chatterbox.mtl_tts import ChatterboxMultilingualTTS  # noqa: E402

VOICES_DIR = os.path.join(ROOT, "models", "voices")
DEFAULT_REF = os.path.join(VOICES_DIR, "referans.wav")
MAX_INPUT_CHARS = 2000
# CHUNK_LIMIT: parca basina uretilecek en uzun metin. KISA tutmak ilk sesi
# hizlandirir (daha az token = daha az otoregresif adim). Iki seviye kullanilir:
#   FIRST_CHUNK_LIMIT : ilk cumle (kullanici en cok bunu bekler)
#   CHUNK_LIMIT       : sonraki cumleler (daha uzun = daha az istek, daha iyi ses)
CHUNK_LIMIT = 220
FIRST_CHUNK_LIMIT = 110
GAP_SECONDS = 0.12         # parcalar arasi sessizlik


def env_float(name, default):
    try:
        return float(os.environ.get(name, default))
    except ValueError:
        return float(default)


class Engine:
    """Modeli ve referans ses onbellegini tutar; uretimi seri yapar (tek GPU)."""

    def __init__(self, device, exaggeration, cfg_weight, temperature, lang):
        self.device = device
        self.exaggeration = exaggeration
        self.cfg_weight = cfg_weight
        self.temperature = temperature
        self.lang = lang
        self.lock = threading.Lock()
        self.conds_cache = {}
        self.compiled = False
        t0 = time.time()
        print("[tts] model yukleniyor (%s)..." % device, flush=True)
        self.model = ChatterboxMultilingualTTS.from_pretrained(device)
        self.sr = self.model.sr
        if self.model.conds is not None:
            self.conds_cache["__yerlesik__"] = self.model.conds
        if device == "cuda":
            # --- GPU hizlandirmalari (olcum raporu: RTF 1.5 -> hedef <0.7) ---
            #
            # 1) TF32: matmul'leri 32-bit floating point tensor cekirdeklerine alir.
            #    Kesinlik kaybi ihmal edilebilir, matmul 2-4x hizlanir.
            #    NOT: torch>=2.1'de varsayilan olarak KAPALI gelir; acikca aciyoruz.
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
            torch.backends.cudnn.benchmark = True  # sabit sekil icin en iyi algoritmayi sec
            try:
                torch.set_float32_matmul_precision("high")
            except Exception:  # noqa: BLE001 - eski torch surumu
                pass
            print("[tts] TF32 + cudnn.benchmark aktif", flush=True)
        print("[tts] model hazir (%.1f sn), ornekleme hizi %d Hz" % (time.time() - t0, self.sr), flush=True)

    def enable_compile(self, mode="reduce-overhead"):
        """T3 dongusunu torch.compile ile derler (CUDA graph -> adim basina 
        baslatma yukunu siler).

        Neden? Chatterbox'in T3 dongusu tamamen GPU'yu bekler; her adimda onlarca
        kucuk kernel baslatilir. CUDA graph bunlari tek bir kayda baglayip tek 
        calistirma indirger; olcum raporunda RTF'in en buyuk tek kalemi bu.

        DIKKAT: Ilk cagri derleme suresi 1-3 dakika HARCAR. Sunucu ilk istegi
        isinma ile karsilarsinca bu gecikme kullaniciya yansimaz. Derleme
        basarisiz olursa sessizce klasik yola donulur (ozelligi kaybetmek, 
        sunucuyu cokertmekten iyidir).
        """
        if self.device != "cuda" or self.compiled:
            return False
        # Derleme hatasi TUM istekleri bozmamalidir. suppress_errors=True ile
        # torch.eager'a sessizce duser (yavas ama calisir).
        # NOT: Bu olmadan derleme hatasi 500 doner ve TTS TAMAMEN calismaz;
        # canli testte boyle bir durum yasandi (Triton yok -> her istek hata).
        try:
            torch._dynamo.config.suppress_errors = True
        except Exception:  # noqa: BLE001 - eski torch
            pass
        try:
            t3 = getattr(self.model, "t3", None)
            if t3 is None or not hasattr(t3, "inference"):
                return False
            t3.inference = torch.compile(t3.inference, mode=mode, fullgraph=False)
            self.compiled = True
            print("[tts] torch.compile(%s) etkin; ilk cagri derlenecek (1-3 dk)"
                  % mode, flush=True)
            return True
        except Exception as e:  # noqa: BLE001 - derleme istege bagli, hata olmali
            print("[tts] UYARI: torch.compile kurulamadi (%r); klasik yol kullanilacak" % (e,),
                  flush=True)
            self.compiled = False
            return False

    def resolve_voice(self, voice):
        """voice alanini bir referans klip yoluna (veya None = yerlesik) cevirir."""
        name = re.sub(r"[^A-Za-z0-9_-]", "", voice or "")
        if name:
            cand = os.path.join(VOICES_DIR, name + ".wav")
            if os.path.isfile(cand):
                return name, cand
        if os.path.isfile(DEFAULT_REF):
            return "referans", DEFAULT_REF
        return "__yerlesik__", None

    def _use_voice(self, key, path):
        cached = self.conds_cache.get(key)
        if cached is None and path:
            self.model.prepare_conditionals(path, exaggeration=self.exaggeration)
            cached = self.model.conds
            self.conds_cache[key] = cached
        if cached is None:
            raise RuntimeError("Yerlesik ses yok; models/voices/referans.wav ekleyin.")
        self.model.conds = cached

    def synth(self, text, voice):
        key, path = self.resolve_voice(voice)
        chunks = split_text(text, CHUNK_LIMIT)
        if not chunks:
            raise ValueError("bos metin")
        # Ilk cumle KISA uretilir: kullanici ilk sesi en cok bekleyen taraftir ve
        # kisa metin = daha az otoregresif adim = daha kisa bekleme.
        # Sonrakiler normal uzunlukta (daha iyi prozonodiy, daha az istek).
        first, rest = chunks[0], chunks[1:]
        if len(first) > FIRST_CHUNK_LIMIT:
            pieces = hard_split(first, FIRST_CHUNK_LIMIT)
            first, rest = pieces[0], pieces[1:] + rest
        ordered = [first] + rest
        parts = []
        gap = np.zeros(int(GAP_SECONDS * self.sr), dtype=np.float32)
        with self.lock, torch.inference_mode():
            self._use_voice(key, path)
            for i, ch in enumerate(ordered):
                wav = self.model.generate(
                    ch,
                    language_id=self.lang,
                    exaggeration=self.exaggeration,
                    cfg_weight=self.cfg_weight,
                    temperature=self.temperature,
                )
                parts.append(wav.squeeze(0).float().cpu().numpy())
                if i < len(ordered) - 1:
                    parts.append(gap)
        return np.concatenate(parts), key


def hard_split(s, limit):
    """limit'ten uzun tek bir cumleyi virgul/bosluk sinirlarindan boler."""
    out, cur = [], ""
    for piece in re.split(r"(?<=[,;:])\s+", s):
        if len(piece) > limit:
            for w in piece.split(" "):
                if len(cur) + len(w) + 1 > limit and cur:
                    out.append(cur)
                    cur = w
                else:
                    cur = (cur + " " + w).strip()
            continue
        if len(cur) + len(piece) + 1 > limit and cur:
            out.append(cur)
            cur = piece
        else:
            cur = (cur + " " + piece).strip()
    if cur:
        out.append(cur)
    return out


def split_text(text, limit):
    text = re.sub(r"\s+", " ", text or "").strip()
    if not text:
        return []
    sentences = re.split(r"(?<=[.!?\u2026])\s+", text)
    chunks, cur = [], ""
    for s in sentences:
        if len(s) > limit:
            if cur:
                chunks.append(cur)
                cur = ""
            chunks.extend(hard_split(s, limit))
        elif len(cur) + len(s) + 1 > limit and cur:
            chunks.append(cur)
            cur = s
        else:
            cur = (cur + " " + s).strip()
    if cur:
        chunks.append(cur)
    return chunks


def to_wav_bytes(samples, sr):
    pcm = np.clip(samples, -1.0, 1.0)
    pcm = (pcm * 32767.0).astype(np.int16)
    buf = io.BytesIO()
    with wave.open(buf, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sr)
        w.writeframes(pcm.tobytes())
    return buf.getvalue()


def create_app(engine):
    app = Flask(__name__)

    @app.get("/health")
    def health():
        _, path = engine.resolve_voice("")
        return jsonify(status="ok", model="chatterbox-multilingual", lang=engine.lang,
                       sample_rate=engine.sr, reference=os.path.basename(path) if path else "yerlesik",
                       exaggeration=engine.exaggeration, cfg_weight=engine.cfg_weight)

    @app.post("/v1/audio/speech")
    def speech():
        body = request.get_json(silent=True) or {}
        text = str(body.get("input", "")).strip()
        if not text:
            return jsonify(error="input bos"), 400
        if len(text) > MAX_INPUT_CHARS:
            return jsonify(error="input cok uzun (en fazla %d karakter)" % MAX_INPUT_CHARS), 413
        t0 = time.time()
        try:
            samples, key = engine.synth(text, str(body.get("voice", "")))
        except Exception as e:  # noqa: BLE001
            print("[tts] HATA: %r" % (e,), flush=True)
            return jsonify(error=str(e)), 500
        dur = len(samples) / float(engine.sr)
        took = time.time() - t0
        print("[tts] %d karakter -> %.1f sn ses, %.1f sn surdu (RTF %.2f), ses=%s"
              % (len(text), dur, took, took / max(dur, 0.01), key), flush=True)
        return Response(to_wav_bytes(samples, engine.sr), mimetype="audio/wav")

    return app


def main():
    ap = argparse.ArgumentParser(description="FIKO AI yerel TTS (Chatterbox)")
    ap.add_argument("port", nargs="?", type=int, default=5000)
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--lang", default=os.environ.get("FIKO_TTS_LANG", "tr"))
    ap.add_argument("--exaggeration", type=float, default=env_float("FIKO_TTS_EXAGGERATION", 0.6))
    ap.add_argument("--cfg", type=float, default=env_float("FIKO_TTS_CFG", 0.3))
    ap.add_argument("--temperature", type=float, default=env_float("FIKO_TTS_TEMPERATURE", 0.8))
    ap.add_argument("--device", default=os.environ.get("FIKO_TTS_DEVICE", ""))
    # --compile: T3 dongusunu torch.compile ile der (CUDA graph). En buyuk RTF
    # kazanci budur. Varsayilan KAPALI (derleme 1-3 dk surer, CPU'da kazanc yok).
    # Kurulum: FIKO_TTS_COMPILE=1 ile de etkinlestirilebilir.
    ap.add_argument("--compile", action="store_true",
                    default=os.environ.get("FIKO_TTS_COMPILE", "") in ("1", "true", "yes", "on"))
    ap.add_argument("--compile-mode", default=os.environ.get("FIKO_TTS_COMPILE_MODE", "reduce-overhead"))
    args = ap.parse_args()

    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
    if device == "cpu":
        print("[tts] UYARI: CUDA yok, CPU'da calisacak (cok yavas).", flush=True)

    engine = Engine(device, args.exaggeration, args.cfg, args.temperature, args.lang)

    # CUDA graph derlemesi: ISINMADAN ONCE yapilir. Yoksa ilk gercek istek
    # 1-3 dakika derleme bekler ve kullanici bunu ilk cumlesinde yasar.
    if args.compile:
        print("[tts] torch.compile etkinlestiriliyor (%s)..." % args.compile_mode, flush=True)
        engine.enable_compile(args.compile_mode)

    # Isinma: ilk gercek istek yavas olmasin.
    try:
        t0 = time.time()
        engine.synth("Merhaba.", "")
        print("[tts] isinma tamam (%.1f sn)" % (time.time() - t0), flush=True)
        if args.compile and engine.compiled:
            # Ikinci bir isinma: derleme gercekten tamamlansin (Ikinci cagri
            # CUDA graph'i gunceller). Sonraki istekler hizli gelsin.
            t0 = time.time()
            engine.synth("Merhaba.", "")
            print("[tts] derleme isinmasi tamam (%.1f sn)" % (time.time() - t0), flush=True)
    except Exception as e:  # noqa: BLE001
        print("[tts] isinma basarisiz: %r" % (e,), flush=True)

    app = create_app(engine)
    _, ref = engine.resolve_voice("")
    print("[tts] dinleniyor http://%s:%d  (referans: %s)" % (args.host, args.port, ref or "yerlesik ses"), flush=True)
    app.run(host=args.host, port=args.port, threaded=True)


if __name__ == "__main__":
    sys.exit(main())