"""voice_daemon.py - resident Cantonese TTS server for LYRA.
Runs in the isolated `envs/voice` environment (qwen-tts, transformers 4.57.3)."""
import json, os, sys, time, traceback
from pathlib import Path
import numpy as np, soundfile as sf, torch
sys.path.insert(0, "/marimo/ai_studio")
from studio_core import BASE, JOBS, JOBS_DONE, RESULTS, LOGS, VOICE_MODEL

MODEL = None
def log(m):
    line = f"[{time.strftime('%H:%M:%S')}] {m}"
    print(line, flush=True)
    with open(LOGS/"voice_daemon.log", "a") as f: f.write(line + "\n")

def get_model():
    global MODEL
    if MODEL is None:
        from qwen_tts import Qwen3TTSModel
        log("load Qwen3-TTS + Cantonese LoRA ...")
        MODEL = Qwen3TTSModel.from_pretrained(str(VOICE_MODEL), device_map="cuda", dtype=torch.bfloat16)
        log(f"voice ready | speakers={MODEL.get_supported_speakers()} | langs={MODEL.get_supported_languages()}")
    return MODEL

def job_tts(j):
    m = get_model()
    spk = j.get("speaker") or (m.get_supported_speakers() or [None])[0]
    kw = dict(text=j["text"], speaker=spk, language=j.get("language") or "chinese")
    if j.get("instruct"): kw["instruct"] = j["instruct"]
    log(f"tts '{j['name']}' spk={spk} :: {j['text'][:30]}")
    t0 = time.time()
    wavs, sr = m.generate_custom_voice(**kw)
    out = BASE/"assets"/"audio"/j["name"]
    sf.write(str(out), np.asarray(wavs[0]), sr)
    dur = len(np.asarray(wavs[0]))/sr
    log(f"wrote {out.name} {dur:.2f}s in {time.time()-t0:.1f}s")
    return {"files": [str(out)], "seconds": round(dur,2), "speaker": spk}

HANDLERS = {"tts": job_tts, "ping": lambda j: {"pong": True}}

def main():
    log(f"voice daemon up pid={os.getpid()}")
    while True:
        found = False
        for f in sorted(JOBS.glob("*.json")):
            try: j = json.loads(f.read_text())
            except Exception: continue
            if j.get("type") not in HANDLERS:     # not ours -> leave it
                continue
            found = True; jid = j.get("id", f.stem)
            try:
                res = {"id": jid, "status": "ok", **(HANDLERS[j["type"]](j) or {})}
            except Exception as e:
                res = {"id": jid, "status": "error", "error": f"{type(e).__name__}: {e}",
                       "trace": traceback.format_exc()[-1200:]}
                log("FAILED " + res["error"])
            tmp = RESULTS/f".{res['id']}.tmp"; tmp.write_text(json.dumps(res, ensure_ascii=False, indent=2))
            tmp.rename(RESULTS/f"{res['id']}.json"); f.rename(JOBS_DONE/f.name)
            log(f"done {res['id']} status={res['status']}")
        if not found: time.sleep(1)

if __name__ == "__main__":
    main()