voice_server.py

import os
import sys

os.environ["COQUI_TOS_AGREED"] = "1"


# ===== Path resolution (works both as script and as PyInstaller onefile) =====
IS_FROZEN = getattr(sys, "frozen", False)
RUNTIME_DIR = os.path.dirname(sys.executable) if IS_FROZEN else os.path.dirname(os.path.abspath(__file__))
BUNDLE_DIR = getattr(sys, "_MEIPASS", RUNTIME_DIR)


def resource_dir(name):
    """Prefer override folders beside the exe/script, then bundled read-only data."""
    runtime_path = os.path.join(RUNTIME_DIR, name)
    bundle_path = os.path.join(BUNDLE_DIR, name)
    if os.path.exists(runtime_path):
        return runtime_path
    return bundle_path


def writable_dir(name):
    return os.path.join(RUNTIME_DIR, name)


BASE_DIR = RUNTIME_DIR
# TTS looks for cached models under $TTS_HOME/tts/... -> anchor to the app dir.
if os.path.exists(os.path.join(BUNDLE_DIR, "tts")) and not os.path.exists(os.path.join(RUNTIME_DIR, "tts")):
    os.environ["TTS_HOME"] = BUNDLE_DIR
else:
    os.environ["TTS_HOME"] = RUNTIME_DIR

import io
import time
import glob
import hashlib
import json
import shutil
import threading
import subprocess

try:
    sys.stdout.reconfigure(line_buffering=True)
except Exception:
    pass


STARTUP_DEBUG = os.environ.get("XTTS_STARTUP_DEBUG", "").strip().lower() in {"1", "true", "yes", "on"}


def startup_log(message):
    if STARTUP_DEBUG:
        print(message, flush=True)


startup_log("[startup] Basic imports complete.")
startup_log("[startup] Importing numpy and torch...")
import numpy as np
import torch
startup_log("[startup] Importing FastAPI and XTTS dependencies...")
from fastapi import FastAPI, HTTPException

# Stream stdout live so the Java client sees [INFO]/[N/4] progress lines as they
# happen (Python block-buffers stdout when it's a pipe, not a console).
try:
    sys.stdout.reconfigure(line_buffering=True)
except Exception:
    pass
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from TTS.api import TTS
import uvicorn
startup_log("[startup] Dependency imports complete.")

AGENTS_DIR = resource_dir("agents")
REFERENCE_CACHE_DIR = writable_dir("reference_cache")
BUNDLED_REFERENCE_CACHE_DIR = os.path.join(BUNDLE_DIR, "reference_cache")
LATENT_CACHE_DIR = writable_dir("latent_cache")
BUNDLED_LATENT_CACHE_DIR = os.path.join(BUNDLE_DIR, "latent_cache")


def _env_int(name, default):
    try:
        return int(os.environ.get(name, default))
    except (TypeError, ValueError):
        return default


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


def _env_bool(name, default):
    value = os.environ.get(name)
    if value is None:
        return default
    return value.strip().lower() not in {"0", "false", "no", "off"}


REFERENCE_SEGMENT_SECONDS = max(6, _env_int("XTTS_REFERENCE_SEGMENT_SECONDS", 12))
REFERENCE_SEGMENT_COUNT = max(1, _env_int("XTTS_REFERENCE_SEGMENT_COUNT", 3))
REFERENCE_MIN_SEGMENT_SECONDS = max(2.0, _env_float("XTTS_REFERENCE_MIN_SECONDS", 3.0))
REFERENCE_SILENCE_THRESHOLD = os.environ.get("XTTS_REFERENCE_SILENCE_DB", "-45dB")
PREPARE_REFERENCES = _env_bool("XTTS_PREPARE_REFERENCES", True)

# Conservative generation defaults trade a little expressiveness for more stable
# speaker likeness and fewer odd XTTS sampling artifacts.
XTTS_DEFAULT_TEMPERATURE = _env_float("XTTS_TEMPERATURE", 0.65)
XTTS_DEFAULT_TOP_P = _env_float("XTTS_TOP_P", 0.80)
XTTS_DEFAULT_TOP_K = _env_int("XTTS_TOP_K", 50)
XTTS_DEFAULT_REPETITION_PENALTY = _env_float("XTTS_REPETITION_PENALTY", 10.0)
XTTS_DEFAULT_LENGTH_PENALTY = _env_float("XTTS_LENGTH_PENALTY", 1.0)
XTTS_MAX_TEXT_CHARS = max(50, _env_int("XTTS_MAX_TEXT_CHARS", 900))
PRELOAD_VOICES = _env_bool("XTTS_PRELOAD_VOICES", False)
PRELOAD_CACHED_VOICES = _env_bool("XTTS_PRELOAD_CACHED_VOICES", True)
WARMUP_INFERENCE = _env_bool("XTTS_WARMUP", True)
LATENT_CACHE_VERSION = 2


def resolve_media_tool(exe_name):
    """Locate ffmpeg/ffprobe: bundled (_MEIPASS), then beside the exe/script
    (in-repo copy), then PATH."""
    cands = []
    if IS_FROZEN:
        cands.append(os.path.join(BUNDLE_DIR, exe_name))
    cands.append(os.path.join(RUNTIME_DIR, exe_name))
    found = shutil.which(exe_name)
    if found:
        cands.append(found)
    for cand in cands:
        if cand and os.path.exists(cand):
            return cand
    return exe_name


FFMPEG = resolve_media_tool("ffmpeg.exe")
FFPROBE = resolve_media_tool("ffprobe.exe")


def _agent_dirs():
    dirs = []
    runtime_agents = os.path.join(RUNTIME_DIR, "agents")
    bundled_agents = os.path.join(BUNDLE_DIR, "agents")
    for path in (runtime_agents, bundled_agents):
        if os.path.isdir(path) and path not in dirs:
            dirs.append(path)
    return dirs


def available_agents():
    agents = set()
    for directory in _agent_dirs():
        for path in glob.glob(os.path.join(directory, "*.mp3")):
            agents.add(os.path.splitext(os.path.basename(path))[0].lower())
    for directory in (LATENT_CACHE_DIR, BUNDLED_LATENT_CACHE_DIR):
        for path in glob.glob(os.path.join(directory, "*.pt")):
            name = os.path.basename(path).split(".", 1)[0].lower()
            if name:
                agents.add(name)
    return sorted(agents)


def find_agent_voice_path(agent):
    filename = f"{agent.lower()}.mp3"
    for directory in _agent_dirs():
        path = os.path.join(directory, filename)
        if os.path.exists(path):
            return path
    return None

# XTTS-v2 in fp32 needs roughly ~2.5-3 GB of VRAM (weights + activations).
# Require this much *free* VRAM before committing to the GPU, otherwise CPU.
XTTS_VRAM_NEED_MIB = 3000

# ===== Initialization Banner =====
print("============================================================")
print("   Valorant Narrator - Voice Server Starting Up...")
print("============================================================")
print("[1/4] Probing hardware and selecting compute device...")


# ===== GPU probe / quick performance test =====
def _query_gpu_util():
    """GPU utilization %% via nvidia-smi. Returns None if unavailable."""
    try:
        out = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=utilization.gpu",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True,
            text=True,
            timeout=3,
        )
        return int(out.stdout.strip().splitlines()[0])
    except Exception:
        return None


def _gpu_benchmark():
    """Small matmul burst to confirm the GPU is alive and responsive.
    Returns (ok, avg_ms_per_iter)."""
    try:
        torch.cuda.synchronize()
        a = torch.randn(2048, 2048, device="cuda")
        b = torch.randn(2048, 2048, device="cuda")
        for _ in range(3):  # warm the kernels
            c = a @ b
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(10):
            c = a @ b
        torch.cuda.synchronize()
        ms = (time.perf_counter() - t0) / 10 * 1000.0
        del a, b, c
        torch.cuda.empty_cache()
        return True, ms
    except Exception as e:
        print(f"[INFO] GPU benchmark failed: {e}")
        return False, 0.0


def select_device():
    if not torch.cuda.is_available():
        print("[INFO] CUDA not available -> using CPU")
        return "cpu"

    try:
        free, total = torch.cuda.mem_get_info()
        free_mib = free // (1024 * 1024)
        total_mib = total // (1024 * 1024)
    except Exception as e:
        print(f"[INFO] Could not query VRAM ({e}) -> using CPU")
        return "cpu"

    name = torch.cuda.get_device_name(0)
    util = _query_gpu_util()
    util_str = f", util {util}%" if util is not None else ""
    print(f"[INFO] GPU: {name} | VRAM free {free_mib}/{total_mib} MiB{util_str}")

    ok, gpu_ms = _gpu_benchmark()
    if not ok:
        print("[INFO] GPU unresponsive -> using CPU")
        return "cpu"
    print(f"[INFO] GPU benchmark: {gpu_ms:.1f} ms / 2048x2048 matmul")

    # Hard constraint: the model must fit with headroom.
    need = XTTS_VRAM_NEED_MIB
    # If the GPU is already busy (e.g. a game is running), demand extra slack.
    if util is not None and util > 85:
        need += 1000
        print(f"[INFO] GPU busy ({util}%) -> raising VRAM requirement to {need} MiB")

    if free_mib < need:
        print(
            f"[INFO] Insufficient free VRAM ({free_mib} MiB < {need} MiB) -> using CPU"
        )
        return "cpu"

    print(f"[INFO] Sufficient VRAM and healthy GPU -> using CUDA")
    return "cuda"


device = select_device()

# ===== Model Initialization =====
print(f"[2/4] Loading XTTS-v2 model on device: {device}")
try:
    print("[INFO] Checking for model files (this may trigger download)...")
    tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2")
    tts.to(device)
except torch.OutOfMemoryError:
    print("[INFO] CUDA OOM while loading -> retrying on CPU")
    torch.cuda.empty_cache()
    device = "cpu"
    tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2")
    tts.to("cpu")
except Exception as e:
    print("[ERROR] Failed to load or download model:", e)
    raise

# Raw XTTS model -> lets us cache speaker latents and call inference() directly,
# skipping the expensive per-request voice-conditioning step.
model = tts.synthesizer.tts_model
model.eval()
torch.set_grad_enabled(False)
# NOTE: do NOT enable cudnn.benchmark — text length varies per request, so every
# call is a new tensor shape and benchmark mode would re-autotune each time.

_cfg = model.config
XTTS_GPT_COND_LEN = max(6, _env_int("XTTS_GPT_COND_LEN", getattr(_cfg, "gpt_cond_len", 30)))
XTTS_GPT_COND_CHUNK_LEN = max(
    2, _env_int("XTTS_GPT_COND_CHUNK_LEN", getattr(_cfg, "gpt_cond_chunk_len", 4))
)
XTTS_GPT_COND_CHUNK_LEN = min(XTTS_GPT_COND_CHUNK_LEN, XTTS_GPT_COND_LEN)
XTTS_MAX_REF_LEN = max(6, _env_int("XTTS_MAX_REF_LEN", getattr(_cfg, "max_ref_len", 30)))
XTTS_SOUND_NORM_REFS = _env_bool(
    "XTTS_SOUND_NORM_REFS", getattr(_cfg, "sound_norm_refs", False)
)
XTTS_LIBROSA_TRIM_DB = _env_int("XTTS_LIBROSA_TRIM_DB", 0)
if XTTS_LIBROSA_TRIM_DB <= 0:
    XTTS_LIBROSA_TRIM_DB = None

print(f"[INFO] ffmpeg: {FFMPEG}")
print(f"[INFO] ffprobe: {FFPROBE}")


def conditioning_kwargs(reference_count):
    if reference_count:
        max_ref_length = REFERENCE_SEGMENT_SECONDS
        total_reference_seconds = max_ref_length * reference_count
    else:
        max_ref_length = XTTS_MAX_REF_LEN
        total_reference_seconds = max_ref_length

    return dict(
        gpt_cond_len=max(1, min(XTTS_GPT_COND_LEN, total_reference_seconds)),
        gpt_cond_chunk_len=XTTS_GPT_COND_CHUNK_LEN,
        max_ref_length=max_ref_length,
        librosa_trim_db=XTTS_LIBROSA_TRIM_DB,
        sound_norm_refs=XTTS_SOUND_NORM_REFS,
    )


def _source_fingerprint(path):
    stat = os.stat(path)
    chunk = 256 * 1024
    digest = hashlib.sha1()
    digest.update(os.path.basename(path).lower().encode("utf-8"))
    digest.update(str(stat.st_size).encode("ascii"))
    with open(path, "rb") as handle:
        digest.update(handle.read(chunk))
        if stat.st_size > chunk:
            handle.seek(max(0, stat.st_size - chunk))
            digest.update(handle.read(chunk))
    return digest.hexdigest()[:12]


def _probe_duration(path):
    try:
        proc = subprocess.run(
            [
                FFPROBE,
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "default=nw=1:nk=1",
                path,
            ],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if proc.returncode != 0:
            return None
        return float(proc.stdout.strip())
    except Exception:
        return None


def _reference_offsets(duration):
    if not duration or duration <= REFERENCE_SEGMENT_SECONDS + 2:
        return [0.0]

    max_start = max(0.0, duration - REFERENCE_SEGMENT_SECONDS - 2)
    if REFERENCE_SEGMENT_COUNT == 1:
        return [0.0]

    fractions = [0.0, 0.33, 0.66, 0.85]
    while len(fractions) < REFERENCE_SEGMENT_COUNT:
        fractions.append(len(fractions) / max(1, REFERENCE_SEGMENT_COUNT))

    offsets = []
    for fraction in fractions[:REFERENCE_SEGMENT_COUNT]:
        offset = min(max_start, max(0.0, duration * fraction))
        if all(abs(offset - existing) > 1.0 for existing in offsets):
            offsets.append(offset)
    return offsets or [0.0]


def _run_reference_ffmpeg(source_path, out_path, offset, audio_filter):
    tmp_path = out_path + ".tmp.wav"
    cmd = [
        FFMPEG,
        "-hide_banner",
        "-loglevel",
        "error",
        "-y",
        "-ss",
        f"{offset:.3f}",
        "-t",
        str(REFERENCE_SEGMENT_SECONDS + 6),
        "-i",
        source_path,
        "-vn",
        "-af",
        audio_filter,
        "-ar",
        "22050",
        "-ac",
        "1",
        "-t",
        str(REFERENCE_SEGMENT_SECONDS),
        tmp_path,
    ]
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        if proc.returncode != 0:
            return False, proc.stderr.strip()
        os.replace(tmp_path, out_path)
        return True, ""
    except Exception as e:
        return False, str(e)
    finally:
        if os.path.exists(tmp_path):
            try:
                os.remove(tmp_path)
            except OSError:
                pass


def _render_reference_segment(source_path, out_path, offset):
    silence = REFERENCE_SILENCE_THRESHOLD
    filters = [
        (
            "highpass=f=80,lowpass=f=9000,"
            f"silenceremove=start_periods=1:start_duration=0.12:start_threshold={silence}:"
            f"stop_periods=-1:stop_duration=0.25:stop_threshold={silence},"
            "loudnorm=I=-23:TP=-2:LRA=11"
        ),
        (
            "highpass=f=80,lowpass=f=9000,"
            f"silenceremove=start_periods=1:start_duration=0.12:start_threshold={silence}:"
            f"stop_periods=-1:stop_duration=0.25:stop_threshold={silence}"
        ),
    ]
    last_error = ""
    for audio_filter in filters:
        ok, last_error = _run_reference_ffmpeg(source_path, out_path, offset, audio_filter)
        if ok:
            return True
    print(f"[INFO] Reference render failed for {os.path.basename(source_path)}: {last_error}")
    return False


def _valid_reference(path):
    if not os.path.exists(path) or os.path.getsize(path) < 4096:
        return False
    duration = _probe_duration(path)
    return duration is None or duration >= REFERENCE_MIN_SEGMENT_SECONDS


def _reference_cache_paths(agent, signature, idx):
    filename = f"{agent}.{signature}.{idx}.wav"
    return [
        os.path.join(REFERENCE_CACHE_DIR, filename),
        os.path.join(BUNDLED_REFERENCE_CACHE_DIR, filename),
    ]


def prepare_reference_audio(agent, voice_path):
    if not PREPARE_REFERENCES:
        return []

    try:
        os.makedirs(REFERENCE_CACHE_DIR, exist_ok=True)
    except OSError as e:
        print(f"[INFO] Could not create reference cache ({e}); using raw references.")
        return []

    try:
        signature = _source_fingerprint(voice_path)
    except OSError as e:
        print(f"[INFO] Could not fingerprint {voice_path} ({e}); using raw reference.")
        return []

    duration = _probe_duration(voice_path)
    offsets = _reference_offsets(duration)
    prepared = []

    for idx, offset in enumerate(offsets, 1):
        runtime_path, bundled_path = _reference_cache_paths(agent, signature, idx)
        existing_path = None
        for candidate in (runtime_path, bundled_path):
            if _valid_reference(candidate):
                existing_path = candidate
                break

        if existing_path is None:
            if not _render_reference_segment(voice_path, runtime_path, offset):
                continue
            existing_path = runtime_path

        if _valid_reference(existing_path):
            prepared.append(existing_path)

    if prepared:
        print(f"[INFO] Prepared {len(prepared)} cleaned reference clip(s) for {agent}")
    return prepared


def _latent_cache_key(agent, source_signature, mode, clips, kwargs):
    payload = {
        "version": LATENT_CACHE_VERSION,
        "agent": agent,
        "source_signature": source_signature,
        "mode": mode,
        "clips": clips,
        "kwargs": kwargs,
        "model": "tts_models/multilingual/multi-dataset/xtts_v2",
    }
    raw = json.dumps(payload, sort_keys=True, separators=(",", ":"))
    return hashlib.sha1(raw.encode("utf-8")).hexdigest()[:16]


def _latent_cache_paths(agent, cache_key):
    filename = f"{agent}.{cache_key}.pt"
    return [
        os.path.join(LATENT_CACHE_DIR, filename),
        os.path.join(BUNDLED_LATENT_CACHE_DIR, filename),
    ]


def _expected_cache_plan(agent, voice_path):
    source_signature = _source_fingerprint(voice_path)
    clips = REFERENCE_SEGMENT_COUNT if PREPARE_REFERENCES else 0
    mode = "cleaned" if clips else "raw"
    kwargs = conditioning_kwargs(clips)
    cache_key = _latent_cache_key(agent, source_signature, mode, clips or 1, kwargs)
    return source_signature, mode, clips, kwargs, cache_key


def load_latents_from_disk(agent, cache_key):
    runtime_path, bundled_path = _latent_cache_paths(agent, cache_key)
    for path, source in ((runtime_path, "runtime"), (bundled_path, "bundled")):
        if not os.path.exists(path):
            continue
        try:
            data = torch.load(path, map_location="cpu")
            if data.get("version") != LATENT_CACHE_VERSION:
                continue
            gpt_cond_latent = data["gpt_cond_latent"].to(device)
            speaker_embedding = data["speaker_embedding"].to(device)
            info = dict(data.get("info") or {})
            info["cache"] = source
            return (gpt_cond_latent, speaker_embedding), info
        except Exception as e:
            print(f"[INFO] Could not load latent cache {path}: {e}")
    return None, None


def load_any_latents_for_agent(agent):
    for directory, source in (
        (LATENT_CACHE_DIR, "runtime"),
        (BUNDLED_LATENT_CACHE_DIR, "bundled"),
    ):
        paths = sorted(
            glob.glob(os.path.join(directory, f"{agent}.*.pt")),
            key=lambda path: os.path.getmtime(path),
            reverse=True,
        )
        for path in paths:
            try:
                data = torch.load(path, map_location="cpu")
                if data.get("version") != LATENT_CACHE_VERSION:
                    continue
                gpt_cond_latent = data["gpt_cond_latent"].to(device)
                speaker_embedding = data["speaker_embedding"].to(device)
                info = dict(data.get("info") or {})
                info["cache"] = source
                info["latent_only"] = True
                return (gpt_cond_latent, speaker_embedding), info
            except Exception as e:
                print(f"[INFO] Could not load latent cache {path}: {e}")
    return None, None


def save_latents_to_disk(agent, cache_key, latents, info):
    try:
        os.makedirs(LATENT_CACHE_DIR, exist_ok=True)
        path = _latent_cache_paths(agent, cache_key)[0]
        tmp_path = path + ".tmp"
        gpt_cond_latent, speaker_embedding = latents
        torch.save(
            {
                "version": LATENT_CACHE_VERSION,
                "agent": agent,
                "gpt_cond_latent": gpt_cond_latent.detach().cpu(),
                "speaker_embedding": speaker_embedding.detach().cpu(),
                "info": info,
            },
            tmp_path,
        )
        os.replace(tmp_path, path)
    except Exception as e:
        print(f"[INFO] Could not save latent cache for {agent}: {e}")

# ===== Per-agent voice-clone cache (pre-allocated conditioning latents) =====
# agent name -> (gpt_cond_latent, speaker_embedding), both already on `device`.
_latent_cache = {}
_reference_info = {}
_cache_lock = threading.Lock()
# Serialize model.inference() — XTTS is not safe for concurrent calls.
_infer_lock = threading.Lock()


def compute_latents(agent):
    """Compute + cache conditioning latents for one agent. Returns the tuple or None."""
    agent = agent.lower()
    cached = _latent_cache.get(agent)
    if cached is not None:
        return cached
    voice_path = find_agent_voice_path(agent)
    if not voice_path:
        cached, info = load_any_latents_for_agent(agent)
        if cached is None:
            return None
        _latent_cache[agent] = cached
        _reference_info[agent] = info
        print(f"[INFO] Loaded latent-only voice for {agent} ({info.get('cache')})")
        return cached
    with _cache_lock:
        cached = _latent_cache.get(agent)
        if cached is not None:
            return cached

        source_signature, mode, clips, kwargs, cache_key = _expected_cache_plan(agent, voice_path)
        cached, info = load_latents_from_disk(agent, cache_key)
        if cached is not None:
            _latent_cache[agent] = cached
            _reference_info[agent] = info
            print(f"[INFO] Loaded cached voice for {agent} ({info.get('cache')})")
            return cached

        reference_paths = prepare_reference_audio(agent, voice_path)
        audio_path = reference_paths if reference_paths else voice_path
        mode = "cleaned" if reference_paths else "raw"
        clips = len(reference_paths) or 1
        kwargs = conditioning_kwargs(len(reference_paths))
        cache_key = _latent_cache_key(agent, source_signature, mode, clips, kwargs)
        cached, info = load_latents_from_disk(agent, cache_key)
        if cached is not None:
            _latent_cache[agent] = cached
            _reference_info[agent] = info
            print(f"[INFO] Loaded cached voice for {agent} ({info.get('cache')})")
            return cached

        try:
            gpt_cond_latent, speaker_embedding = model.get_conditioning_latents(
                audio_path=audio_path, **kwargs
            )
            info = {
                "mode": mode,
                "clips": len(reference_paths) or 1,
                "gpt_cond_len": kwargs["gpt_cond_len"],
                "cache": "generated",
            }
        except Exception as e:
            if not reference_paths:
                raise
            print(f"[INFO] Cleaned references failed for {agent}: {e}; retrying raw audio")
            kwargs = conditioning_kwargs(0)
            mode = "raw-fallback"
            clips = 1
            cache_key = _latent_cache_key(agent, source_signature, mode, clips, kwargs)
            cached, info = load_latents_from_disk(agent, cache_key)
            if cached is not None:
                _latent_cache[agent] = cached
                _reference_info[agent] = info
                print(f"[INFO] Loaded cached voice for {agent} ({info.get('cache')})")
                return cached
            gpt_cond_latent, speaker_embedding = model.get_conditioning_latents(
                audio_path=voice_path, **kwargs
            )
            info = {
                "mode": mode,
                "clips": 1,
                "gpt_cond_len": kwargs["gpt_cond_len"],
                "cache": "generated",
            }

        _latent_cache[agent] = (gpt_cond_latent, speaker_embedding)
        _reference_info[agent] = info
        save_latents_to_disk(agent, cache_key, _latent_cache[agent], info)
        return _latent_cache[agent]


print("[3/4] Discovering agent voices...")
_available_agents = available_agents()
if not _available_agents:
    print(f"[INFO] No agent voices found. Checked: {_agent_dirs()}")
else:
    print(f"[INFO] Found {len(_available_agents)} agent voice(s).")

if PRELOAD_CACHED_VOICES and _available_agents:
    loaded = 0
    for _agent in _available_agents:
        try:
            _voice_path = find_agent_voice_path(_agent)
            if _voice_path:
                _, _, _, _, _cache_key = _expected_cache_plan(_agent, _voice_path)
                _cached, _info = load_latents_from_disk(_agent, _cache_key)
            else:
                _cached, _info = load_any_latents_for_agent(_agent)
            if _cached is None:
                continue
            _latent_cache[_agent] = _cached
            _reference_info[_agent] = _info
            loaded += 1
        except Exception as e:
            print(f"[INFO] Could not preload cached voice {_agent}: {e}")
    print(f"[INFO] Loaded {loaded}/{len(_available_agents)} cached voice clone(s) from disk.")

if PRELOAD_VOICES and _available_agents:
    print("[INFO] XTTS_PRELOAD_VOICES=1 -> precomputing all voice clones before startup.")
    for _i, _agent in enumerate(_available_agents, 1):
        try:
            compute_latents(_agent)
            print(f"[INFO] Cached voice {_i}/{len(_available_agents)}: {_agent}")
        except Exception as e:
            print(f"[INFO] Could not cache {_agent}: {e}")
else:
    print("[INFO] Voice clones will be prepared lazily on first use.")

# ===== Warmup — compile CUDA kernels so the first real request is fast =====
print("[4/4] Warming up inference engine...")
try:
    if WARMUP_INFERENCE and _latent_cache:
        _g, _s = next(iter(_latent_cache.values()))
        with _infer_lock:
            # Mirror the real request path (enable_text_splitting=True) so the
            # sentence splitter and all CUDA kernels are initialized up front.
            for _wtext in ("Warming up the engine now.", "Ready."):
                model.inference(_wtext, "en", _g, _s, enable_text_splitting=True)
        if device == "cuda":
            torch.cuda.synchronize()
        print("[INFO] Warmup complete.")
    elif WARMUP_INFERENCE:
        print("[INFO] Warmup skipped: no cached voices are loaded yet.")
    else:
        print("[INFO] Warmup disabled for faster startup.")
except Exception as e:
    print(f"[INFO] Warmup skipped: {e}")


def encode_mp3(wav, sr):
    """Encode a float32 mono waveform to 44.1kHz stereo 192k mp3, in memory."""
    pcm = np.asarray(wav, dtype=np.float32).tobytes()
    cmd = [
        FFMPEG, "-hide_banner", "-loglevel", "error", "-y",
        "-f", "f32le", "-ar", str(sr), "-ac", "1", "-i", "pipe:0",
        "-vn", "-ar", "44100", "-ac", "2", "-b:a", "192k", "-f", "mp3", "pipe:1",
    ]
    proc = subprocess.run(
        cmd, input=pcm, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    if proc.returncode != 0 or not proc.stdout:
        error = proc.stderr.decode("utf-8", errors="ignore").strip()
        raise RuntimeError(f"ffmpeg mp3 encoding failed: {error or proc.returncode}")
    return proc.stdout


print("============================================================")
print("Application startup complete - Voice server is ready.")
print("============================================================\n")

# ===== FastAPI Setup =====
app = FastAPI()
SUPPORTED_LANGUAGES = {str(lang).lower() for lang in getattr(_cfg, "languages", [])}


def _clamp_float(value, default, minimum, maximum):
    try:
        value = default if value is None else float(value)
    except (TypeError, ValueError):
        value = default
    return max(minimum, min(maximum, value))


def _clamp_int(value, default, minimum, maximum):
    try:
        value = default if value is None else int(value)
    except (TypeError, ValueError):
        value = default
    return max(minimum, min(maximum, value))


def _normalize_language(language):
    lang = (language or "en").strip().lower().replace("_", "-")
    if lang in SUPPORTED_LANGUAGES:
        return lang
    base = lang.split("-")[0]
    if base in SUPPORTED_LANGUAGES:
        return base
    if lang == "zh" and "zh-cn" in SUPPORTED_LANGUAGES:
        return "zh-cn"
    return "en"


def _generation_kwargs(req):
    return {
        "temperature": _clamp_float(req.temperature, XTTS_DEFAULT_TEMPERATURE, 0.45, 0.95),
        "top_p": _clamp_float(req.top_p, XTTS_DEFAULT_TOP_P, 0.60, 0.95),
        "top_k": _clamp_int(req.top_k, XTTS_DEFAULT_TOP_K, 10, 100),
        "repetition_penalty": _clamp_float(
            req.repetition_penalty, XTTS_DEFAULT_REPETITION_PENALTY, 2.0, 12.0
        ),
        "length_penalty": _clamp_float(
            req.length_penalty, XTTS_DEFAULT_LENGTH_PENALTY, 0.5, 2.0
        ),
        "speed": _clamp_float(req.speed, 1.0, 0.75, 1.25),
    }


class SpeakRequest(BaseModel):
    agent: str
    text: str
    language: str = "en"
    speed: float | None = None
    temperature: float | None = None
    top_p: float | None = None
    top_k: int | None = None
    repetition_penalty: float | None = None
    length_penalty: float | None = None


@app.get("/health")
def health():
    info = {
        "status": "ok",
        "device": device,
        "agents": _available_agents,
        "cached_agents": sorted(_latent_cache.keys()),
        "lazy_voice_loading": not PRELOAD_VOICES,
        "reference_mode": "cleaned" if PREPARE_REFERENCES else "raw",
        "reference_cache_dir": REFERENCE_CACHE_DIR,
        "latent_cache_dir": LATENT_CACHE_DIR,
        "references": _reference_info,
        "conditioning": {
            "gpt_cond_len": XTTS_GPT_COND_LEN,
            "gpt_cond_chunk_len": XTTS_GPT_COND_CHUNK_LEN,
            "max_ref_len": XTTS_MAX_REF_LEN,
            "librosa_trim_db": XTTS_LIBROSA_TRIM_DB,
            "sound_norm_refs": XTTS_SOUND_NORM_REFS,
        },
        "generation_defaults": {
            "temperature": XTTS_DEFAULT_TEMPERATURE,
            "top_p": XTTS_DEFAULT_TOP_P,
            "top_k": XTTS_DEFAULT_TOP_K,
            "repetition_penalty": XTTS_DEFAULT_REPETITION_PENALTY,
            "length_penalty": XTTS_DEFAULT_LENGTH_PENALTY,
            "max_text_chars": XTTS_MAX_TEXT_CHARS,
        },
    }
    if device == "cuda":
        try:
            free, total = torch.cuda.mem_get_info()
            info["vram_free_mib"] = free // (1024 * 1024)
            info["vram_total_mib"] = total // (1024 * 1024)
        except Exception:
            pass
    return JSONResponse(info)


@app.post("/speak")
def speak(req: SpeakRequest):
    agent_name = (req.agent or "").strip().lower()
    text = (req.text or "").strip()
    if not agent_name:
        raise HTTPException(status_code=400, detail="agent is required")
    if not text:
        raise HTTPException(status_code=400, detail="text is required")
    if len(text) > XTTS_MAX_TEXT_CHARS:
        raise HTTPException(
            status_code=400,
            detail=f"text is too long ({len(text)} > {XTTS_MAX_TEXT_CHARS} chars)",
        )

    latents = compute_latents(agent_name)  # cached hit, or lazily compute+cache
    if latents is None:
        raise HTTPException(
            status_code=404, detail=f"Voice file not found for agent: {req.agent}"
        )
    gpt_cond_latent, speaker_embedding = latents
    language = _normalize_language(req.language)
    generation = _generation_kwargs(req)

    t0 = time.perf_counter()
    with _infer_lock:
        out = model.inference(
            text,
            language,
            gpt_cond_latent,
            speaker_embedding,
            enable_text_splitting=True,
            **generation,
        )
        if device == "cuda":
            torch.cuda.synchronize()
    infer_ms = (time.perf_counter() - t0) * 1000.0

    wav = out["wav"] if isinstance(out, dict) else out
    sr = getattr(model.config, "output_sample_rate", 24000)
    mp3_bytes = encode_mp3(wav, sr)

    total_ms = (time.perf_counter() - t0) * 1000.0
    print(
        f"[INFO] {agent_name}: {len(text)} chars | lang {language} | "
        f"temp {generation['temperature']:.2f} | speed {generation['speed']:.2f} | "
        f"infer {infer_ms:.0f} ms | total {total_ms:.0f} ms"
    )

    return StreamingResponse(io.BytesIO(mp3_bytes), media_type="audio/mpeg")


if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=5005)