#!/usr/bin/env python3
# Blend Kokoro voicepacks into a new persistent voice.
#
# Mirrors Kokoro-FastAPI's request-time blend op (see
# api/src/services/tts_service.py::_get_voices_path) so the resulting
# .pt behaves identically to the equivalent inline `voice="a(w)+b(w)"`
# request — except it lives on disk, gets a stable name, and is
# discoverable via GET /v1/audio/voices.
#
# Designed to be run INSIDE the kokoro container (where torch and the
# built-in voicepacks are already present):
#
#   docker exec -u 0 kokoro python /tmp/blend_kokoro_voice.py \
#     --recipe 'af_bella(2)+af_aoede(1)' --out athena
#
# I/O layout (paired with the compose changes that ship custom-voice
# persistence):
#
#   * Reads from /app/user_voices/, which the compose `command:` shim
#     pre-seeds with the in-image built-ins on every container start.
#     So both built-ins (af_bella, af_aoede, ...) and any prior custom
#     blends resolve from the same directory.
#   * Writes to /app/user_voices/<out>.pt — that path is the host
#     bind mount /worktank/kokoro/user_voices/, so the result is
#     immediately discoverable (the wrapper has VOICES_DIR pointed
#     here) AND fully persistent across container recreate, image
#     upgrade, and host reboot.
import argparse
import os
import re
import sys

import torch

VOICES_DIR = "/app/user_voices"


def resolve(name: str) -> str:
    p = os.path.join(VOICES_DIR, f"{name}.pt")
    if os.path.exists(p):
        return p
    raise FileNotFoundError(f"voice {name!r} not found in {VOICES_DIR}")


def parse_recipe(recipe: str):
    parts = re.split(r"([-+])", recipe.replace(" ", ""))
    terms = []
    for i, tok in enumerate(parts):
        if i % 2 == 0:
            if "(" in tok and ")" in tok:
                name = tok.split("(")[0]
                weight = float(tok.split("(")[1].split(")")[0])
            else:
                name, weight = tok, 1.0
            sign = +1.0 if i == 0 or parts[i - 1] == "+" else -1.0
            terms.append((name, sign * weight))
    return terms


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--recipe", required=True,
                    help='e.g. "af_bella(2)+af_aoede(1)"')
    ap.add_argument("--out", required=True, help="output voice name (no .pt)")
    ap.add_argument("--no-normalize", action="store_true",
                    help="sum weights as-is instead of dividing by total")
    ap.add_argument("--force", action="store_true",
                    help="overwrite existing output voice")
    args = ap.parse_args()

    terms = parse_recipe(args.recipe)
    if not terms:
        sys.exit("empty recipe")

    total = sum(abs(w) for _, w in terms) if not args.no_normalize else 1.0
    if total == 0:
        sys.exit("weights sum to zero")

    out_path = os.path.join(VOICES_DIR, f"{args.out}.pt")
    if os.path.exists(out_path) and not args.force:
        sys.exit(f"refusing to overwrite {out_path} (pass --force)")

    blended = None
    for name, w in terms:
        t = torch.load(resolve(name), map_location="cpu") * (w / total)
        blended = t if blended is None else blended + t

    torch.save(blended, out_path)
    print(f"wrote {out_path} (shape={tuple(blended.shape)}, dtype={blended.dtype})")


if __name__ == "__main__":
    main()