from __future__ import annotations

import argparse
import asyncio
import json
import os
import struct
import sys
import tempfile
from pathlib import Path
from typing import Any

import numpy as np

MODEL_SAMPLE_RATE = 24000
OUTPUT_SAMPLE_RATE = 48000
STREAM_CHUNK_MS = 80
STREAMING_INTERVAL = 0.12
CHUNK_BYTES = OUTPUT_SAMPLE_RATE * STREAM_CHUNK_MS // 1000 * 2  # 7680


def load_config(path: str) -> dict[str, Any]:
    with open(path, "r", encoding="utf-8") as f:
        cfg = json.load(f)
    cfg["_config_dir"] = str(Path(path).expanduser().resolve().parent)
    return cfg


def voice_by_id(cfg: dict[str, Any], voice_id: str) -> dict[str, Any]:
    for voice in cfg.get("voices", []):
        if voice.get("id") == voice_id:
            return voice
    raise SystemExit(f"voice not found: {voice_id}")


def resolve_path(cfg: dict[str, Any], value: str) -> str:
    p = Path(value).expanduser()
    if p.is_absolute():
        return str(p)
    return str(Path(cfg["_config_dir"]) / p)


def preprocess_ref_audio(path: str) -> str:
    try:
        import librosa
        import soundfile as sf

        audio, _ = librosa.load(path, sr=MODEL_SAMPLE_RATE, mono=True)
        audio = audio - np.mean(audio)
        audio, _ = librosa.effects.trim(audio, top_db=40)
        audio = np.pad(audio, (0, int(0.1 * MODEL_SAMPLE_RATE)))
        peak = np.max(np.abs(audio))
        if peak > 0.95:
            audio = audio * 0.95 / peak
        tmp_path = os.path.join(tempfile.gettempdir(), f"ivox_ref_{hash(path)}.wav")
        sf.write(tmp_path, audio, MODEL_SAMPLE_RATE)
        return tmp_path
    except Exception as e:
        print(f"reference preprocess failed: {e}", file=sys.stderr)
        return path


def array_to_pcm(audio: np.ndarray, peak_limit: float = 0.9) -> bytes:
    if peak_limit > 0 and audio.size:
        peak = np.max(np.abs(audio))
        if peak > peak_limit:
            audio = audio * (peak_limit / peak)
    audio = np.clip(audio * 32767, -32768, 32767).astype(np.int16)
    return audio.tobytes()


def resample(pcm: bytes, from_rate: int, to_rate: int) -> bytes:
    if from_rate == to_rate:
        return pcm
    samples = np.frombuffer(pcm, dtype=np.int16)
    if to_rate == from_rate * 2:
        return np.repeat(samples, 2).tobytes()
    from scipy import signal as scipy_signal

    samples_f32 = samples.astype(np.float32)
    target_len = int(len(samples_f32) * to_rate / from_rate)
    out = scipy_signal.resample(samples_f32, target_len)
    return np.clip(np.round(out), -32768, 32767).astype(np.int16).tobytes()


# ── binary framing helpers ──


def _write_i64(n: int) -> None:
    sys.stdout.buffer.write(struct.pack("<q", n))


def _write_frame(data: bytes) -> None:
    _write_i64(len(data))
    sys.stdout.buffer.write(data)
    sys.stdout.buffer.flush()


def _send_ok() -> None:
    _write_i64(0)
    sys.stdout.buffer.flush()


def _send_chunk(pcm: bytes) -> None:
    _write_frame(pcm)


def _send_end() -> None:
    _write_i64(0)
    sys.stdout.buffer.flush()


def _send_error(msg: str) -> None:
    _write_i64(1)
    _write_frame(msg.encode("utf-8"))


# ── one-shot synthesize (backward compat) ──


async def synthesize(args: argparse.Namespace) -> None:
    cfg = load_config(args.config)
    voice = voice_by_id(cfg, args.voice)
    model_path = cfg.get("model", {}).get("path")
    if not model_path:
        raise SystemExit("model.path is required")

    from mlx_audio.tts.utils import load_model

    model = load_model(model_path)
    ref_audio = voice.get("refAudio") or ""
    ref_text = voice.get("refText") or ""
    if ref_audio:
        ref_audio = preprocess_ref_audio(resolve_path(cfg, ref_audio))

    try:
        stream = model.generate(
            text=args.text,
            ref_audio=ref_audio or None,
            ref_text=ref_text or None,
            stream=True,
            streaming_interval=STREAMING_INTERVAL,
            lang_code="chinese",
        )
    except TypeError:
        stream = model.generate(text=args.text, stream=True, lang_code="chinese")

    for result in stream:
        audio = np.array(result.audio)
        if audio.size == 0:
            continue
        pcm = resample(array_to_pcm(audio), MODEL_SAMPLE_RATE, OUTPUT_SAMPLE_RATE)
        for offset in range(0, len(pcm), CHUNK_BYTES):
            sys.stdout.buffer.write(pcm[offset : offset + CHUNK_BYTES])
            sys.stdout.buffer.flush()


# ── daemon mode (常驻进程) ──


def _handle_request(cfg: dict[str, Any], model: Any, req: dict[str, Any]) -> None:
    voice_id = req.get("voice_id", "")
    text = req.get("text", "")
    if not text or not voice_id:
        _send_error("text and voice_id required")
        return

    voice = voice_by_id(cfg, voice_id)
    ref_audio = voice.get("refAudio") or ""
    ref_text = voice.get("refText") or ""
    if ref_audio:
        ref_audio = preprocess_ref_audio(resolve_path(cfg, ref_audio))

    try:
        stream = model.generate(
            text=text,
            ref_audio=ref_audio or None,
            ref_text=ref_text or None,
            stream=True,
            streaming_interval=STREAMING_INTERVAL,
            lang_code="chinese",
        )
    except TypeError:
        stream = model.generate(text=text, stream=True, lang_code="chinese")

    _send_ok()
    try:
        for result in stream:
            audio = np.array(result.audio)
            if audio.size == 0:
                continue
            pcm = resample(array_to_pcm(audio), MODEL_SAMPLE_RATE, OUTPUT_SAMPLE_RATE)
            for offset in range(0, len(pcm), CHUNK_BYTES):
                _send_chunk(pcm[offset : offset + CHUNK_BYTES])
    except Exception as e:
        print(f"generate error: {e}", file=sys.stderr)
    _send_end()


def run_daemon(args: argparse.Namespace) -> None:
    cfg = load_config(args.config)
    model_path = cfg.get("model", {}).get("path")
    if not model_path:
        _send_error("model.path is required")
        sys.exit(1)

    # 预热：预导入 librosa 避免首次请求卡顿
    try:
        import librosa
        import soundfile as sf
    except Exception:
        pass

    from mlx_audio.tts.utils import load_model

    model = load_model(model_path)

    # 预热模型推理（消除首次 JIT 编译延迟）
    default_voice = cfg.get("voices", [{}])[0]
    ref_audio = default_voice.get("refAudio") or ""
    ref_text = default_voice.get("refText") or ""
    if ref_audio:
        ref_audio = preprocess_ref_audio(resolve_path(cfg, ref_audio))
    try:
        stream = model.generate(
            text=".",
            ref_audio=ref_audio or None,
            ref_text=ref_text or None,
            stream=True,
            streaming_interval=STREAMING_INTERVAL,
            lang_code="chinese",
        )
        for result in stream:
            pass  # 消耗完预热流
    except Exception:
        pass  # 预热失败不影响启动

    # 通知 Go 模型就绪
    sys.stdout.write("READY\n")
    sys.stdout.flush()

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
        except json.JSONDecodeError as e:
            _send_error(f"json decode: {e}")
            continue
        _handle_request(cfg, model, req)


# ── CLI ──


def list_voices(args: argparse.Namespace) -> None:
    cfg = load_config(args.config)
    for voice in cfg.get("voices", []):
        print(f"{voice.get('id')}\t{voice.get('name', '')}\t{voice.get('description', '')}")


def main() -> None:
    parser = argparse.ArgumentParser(prog="ivox_tts.worker")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("synthesize")
    p.add_argument("--config", required=True)
    p.add_argument("--voice", required=True)
    p.add_argument("--text", required=True)

    p = sub.add_parser("daemon")
    p.add_argument("--config", required=True)

    p = sub.add_parser("voices")
    p.add_argument("--config", required=True)

    args = parser.parse_args()
    if args.cmd == "synthesize":
        asyncio.run(synthesize(args))
    elif args.cmd == "daemon":
        run_daemon(args)
    elif args.cmd == "voices":
        list_voices(args)


if __name__ == "__main__":
    main()
