#!/usr/bin/env python3
"""One-command local text to speech (Kokoro) from GenLovers.

Usage:
    python setup_text_to_speech.py "Hello, this is Kokoro speaking."
    python setup_text_to_speech.py "Bonjour" --lang fr-fr --voice ff_siwis --out bonjour.wav

First run: creates a private virtual environment, installs kokoro-onnx into it,
and downloads the Kokoro model files (about 325 MB). Later runs skip both. The
speech is written to a WAV file (speech.wav by default). Everything runs on this
machine on CPU; nothing is uploaded anywhere. Needs Python 3.10 or newer (tested on 3.13).
"""
import argparse
import os
import subprocess
import sys
import urllib.request
import venv
from pathlib import Path

ROOT = Path(os.environ.get("GENLOVERS_TOOLS_DIR", Path.home() / "genlovers-tools"))
ENV_DIR = ROOT / "text-to-speech"
BIN = ENV_DIR / ("Scripts" if os.name == "nt" else "bin")
PY = BIN / ("python.exe" if os.name == "nt" else "python")
MODEL_BASE = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0"
FILES = ["kokoro-v1.0.onnx", "voices-v1.0.bin"]
PACKAGES = ["kokoro-onnx", "soundfile"]

# Runs inside the private environment, so the packages are importable there.
RUNNER = """
import sys
import soundfile as sf
from kokoro_onnx import Kokoro
model, voices, text, voice, lang, speed, out = sys.argv[1:8]
kokoro = Kokoro(model, voices)
samples, rate = kokoro.create(text, voice=voice, speed=float(speed), lang=lang)
sf.write(out, samples, rate)
"""


def say(msg):
    print(f"\n==> {msg}", flush=True)


def fail(msg):
    print(f"\nSETUP FAILED: {msg}", file=sys.stderr)
    sys.exit(1)


def ensure_installed():
    if sys.version_info < (3, 10):
        fail("Python 3.10 or newer is required. Install it from https://www.python.org/downloads/")
    if not PY.exists():
        say(f"Creating a private environment in {ENV_DIR}")
        venv.create(ENV_DIR, with_pip=True)
    check = subprocess.run([str(PY), "-c", "import kokoro_onnx, soundfile"], capture_output=True)
    if check.returncode != 0:
        say("Installing kokoro-onnx (this can take a few minutes)...")
        if subprocess.run([str(PY), "-m", "pip", "install", "--upgrade", *PACKAGES]).returncode != 0:
            fail("pip could not install kokoro-onnx. Try Python 3.13, the version this script was tested on, and run this script again.")
    for name in FILES:
        target = ENV_DIR / name
        if target.exists():
            continue
        say(f"Downloading {name}...")
        partial = target.with_suffix(target.suffix + ".part")
        try:
            urllib.request.urlretrieve(f"{MODEL_BASE}/{name}", partial)
        except OSError as err:
            fail(f"Could not download {name}: {err}")
        partial.replace(target)


def main():
    parser = argparse.ArgumentParser(description="Local Kokoro text to speech.")
    parser.add_argument("text", nargs="?", help="The line to speak.")
    parser.add_argument("--voice", default="af_heart", help="Voice name (default af_heart).")
    parser.add_argument("--lang", default="en-us", help="Language code: en-us, en-gb, fr-fr, es, it, pt-br or hi (default en-us).")
    parser.add_argument("--speed", default="1.0", help="Speaking speed from 0.5 to 2.0 (default 1.0).")
    parser.add_argument("--out", default="speech.wav", help="Output WAV file (default speech.wav).")
    args = parser.parse_args()

    ensure_installed()
    if not args.text:
        print(
            "\nSETUP COMPLETE\nNow generate speech with:\n"
            '    python setup_text_to_speech.py "Hello from my own machine"'
        )
        return
    say("Generating speech...")
    result = subprocess.run(
        [
            str(PY), "-c", RUNNER,
            str(ENV_DIR / FILES[0]), str(ENV_DIR / FILES[1]),
            args.text, args.voice, args.lang, args.speed, args.out,
        ]
    )
    if result.returncode != 0:
        fail(
            "Kokoro could not synthesize. On Linux, install the espeak-ng package "
            "(for example: sudo apt install espeak-ng) and try again."
        )
    print(f"Saved {args.out}")


if __name__ == "__main__":
    main()