#!/usr/bin/env python3
"""Registra le voci clonate (.spk/.rvq) trovate in VOICE_DIR nel tts-server.

Container-friendly: scansiona VOICE_DIR ricorsivamente; per ogni *.spk
cerca il .rvq omonimo (obbligatorio) e il .txt omonimo (ref_text, opzionale;
fallback: *_batch_ref.txt nella stessa directory).
"""
import base64
import json
import os
import sys
import urllib.request

PORT = os.environ.get("TTS_PORT", "8881")
URL = f"http://127.0.0.1:{PORT}/v1/audio/voices"
VOICE_DIR = os.environ.get("VOICE_DIR", "/voices")


def register(name, spk_path, rvq_path, ref_text):
    spk_b64 = base64.b64encode(open(spk_path, "rb").read()).decode()
    rvq_b64 = base64.b64encode(open(rvq_path, "rb").read()).decode()
    body = {"name": name, "spk_b64": spk_b64, "rvq_b64": rvq_b64}
    if ref_text:
        body["ref_text"] = ref_text
    req = urllib.request.Request(
        URL, data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"}, method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as r:
        resp = json.loads(r.read().decode())
        print(f"{name}: {r.status} {resp.get('status')}")
        return True


def main():
    if not os.path.isdir(VOICE_DIR):
        print(f"VOICE_DIR {VOICE_DIR} non esiste, salto registrazione")
        return 0
    ok = total = 0
    for root, _dirs, files in os.walk(VOICE_DIR):
        batch_ref = None
        for f in files:
            if f.endswith("_batch_ref.txt"):
                batch_ref = open(os.path.join(root, f)).read()
        for f in sorted(files):
            if not f.endswith(".spk"):
                continue
            stem = f[: -len(".spk")]
            name = stem.replace("_ref", "").replace("_vocals", "")
            spk_path = os.path.join(root, f)
            rvq_path = os.path.join(root, stem + ".rvq")
            if not os.path.isfile(rvq_path):
                print(f"{name}: rvq mancante ({rvq_path}), salto")
                continue
            txt_path = os.path.join(root, stem + ".txt")
            ref_text = None
            if os.path.isfile(txt_path):
                ref_text = open(txt_path).read()
            elif batch_ref:
                ref_text = batch_ref
            total += 1
            try:
                if register(name, spk_path, rvq_path, ref_text):
                    ok += 1
            except Exception as e:
                print(f"{name}: ERRORE {e}")
    print(f"\nRegistrate {ok}/{total} voci")
    return 0 if ok == total else 1


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