"""Qwen3-TTS HTTP server — OpenAI-compatible speech endpoint with voice training."""
import argparse, io, json, os, re, sys, tempfile, time
from pathlib import Path
from flask import Flask, request, send_file, jsonify

app = Flask(__name__)
MODEL_DIR = "/models/qwen3-tts"
CUSTOM_VOICES_DIR = "/models/qwen3-tts/custom_speakers"
engine = None

PRESET_SPEAKERS = ["Vivian","Serena","uncle_fu","Dylan","Eric","Ryan","Aiden","ono_anna","Sohee"]
LANGUAGES = ["english","chinese","japanese","korean","german","spanish","french","russian","italian","portuguese"]
SAFE_NAME_RE = re.compile(r'^[a-zA-Z0-9_-]{1,64}$')

# ── Bootstrap helpers ──

def ensure_embeddings(model_dir):
    emb_dir = os.path.join(model_dir, "embeddings")
    if os.path.exists(os.path.join(emb_dir, "text_embedding_projected.npy")):
        return
    gguf_path = os.path.join(model_dir, "qwen3_assets.gguf")
    if not os.path.exists(gguf_path):
        print(f"[server] WARNING: {gguf_path} not found", flush=True)
        return
    print(f"[server] Extracting embeddings from GGUF...", flush=True)
    import numpy as np
    from gguf import GGUFReader
    NAME_MAP = {"text_embd": "text_embedding_projected.npy",
                "proj.weight": "proj_weight.npy", "proj.bias": "proj_bias.npy"}
    for j in range(16):
        NAME_MAP[f"codec_embd.{j}"] = f"codec_embedding_{j}.npy"
    reader = GGUFReader(gguf_path)
    os.makedirs(emb_dir, exist_ok=True)
    for tensor in reader.tensors:
        outname = NAME_MAP.get(tensor.name)
        if outname is None:
            continue
        data = tensor.data
        shape = list(tensor.shape)
        n_el = 1
        for d in shape:
            n_el *= d
        n_blocks = n_el // 32
        raw = data.tobytes()
        block_sz = 34
        if len(raw) == n_blocks * block_sz:
            quants = np.zeros(n_el, dtype=np.float32)
            for bi in range(n_blocks):
                off = bi * block_sz
                s = float(np.frombuffer(raw[off:off+2], dtype=np.float16)[0])
                qs = np.frombuffer(raw[off+2:off+block_sz], dtype=np.int8)
                quants[bi*32:(bi+1)*32] = qs.astype(np.float32) * s
            arr = quants.reshape(shape)
        else:
            arr = np.array(data, dtype=np.float32).reshape(shape)
        np.save(os.path.join(emb_dir, outname), arr)
    print(f"[server] Embeddings extracted", flush=True)

def ensure_symlinks(model_dir):
    links = {"qwen3_tts_talker.q5_k.gguf": "qwen3_tts_talker.gguf",
             "qwen3_tts_predictor.q8_0.gguf": "qwen3_tts_predictor.gguf",
             "qwen3_tts_decoder.fp16.onnx": "qwen3_tts_decoder.onnx",
             "qwen3_tts_codec_encoder.fp16.onnx": "qwen3_tts_codec_encoder.onnx",
             "qwen3_tts_speaker_encoder.fp16.onnx": "qwen3_tts_speaker_encoder.onnx"}
    for link_name, target in links.items():
        lp = os.path.join(model_dir, link_name)
        tp = os.path.join(model_dir, target)
        if not os.path.exists(lp) and os.path.exists(tp):
            os.symlink(target, lp)

def init_engine():
    global engine
    parent = os.path.dirname(os.path.abspath(MODEL_DIR))
    basename = os.path.basename(MODEL_DIR)
    ensure_embeddings(MODEL_DIR)
    ensure_symlinks(MODEL_DIR)
    os.makedirs(CUSTOM_VOICES_DIR, exist_ok=True)
    os.chdir(parent)
    from qwen3_tts_gguf.inference import TTSEngine
    engine = TTSEngine(model_dir=basename, onnx_provider="CPUExecutionProvider")
    if not engine.ready:
        print(f"[server] WARNING: engine created but not ready", flush=True)

# ── Voice management helpers ──

def list_all_speakers():
    presets = []
    for name in PRESET_SPEAKERS:
        presets.append({"name": name, "type": "preset"})
    customs = []
    if os.path.isdir(CUSTOM_VOICES_DIR):
        for f in sorted(os.listdir(CUSTOM_VOICES_DIR)):
            if f.endswith(".json"):
                vname = f[:-5]
                meta = _load_voice_meta(vname)
                customs.append({"name": vname, "type": "custom",
                                "description": meta.get("description", ""),
                                "created": meta.get("created", "")})
    return presets + customs

def _load_voice_meta(name):
    path = os.path.join(CUSTOM_VOICES_DIR, f"{name}.json")
    if not os.path.exists(path):
        return {}
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    return {"description": data.get("description", ""),
            "created": data.get("created", ""),
            "ref_text": data.get("text", ""),
            "duration_hint": data.get("duration_hint", 0)}

def _voice_exists(name):
    return os.path.exists(os.path.join(CUSTOM_VOICES_DIR, f"{name}.json"))

def _is_custom_voice(name):
    return name not in [s.lower() for s in PRESET_SPEAKERS] and _voice_exists(name)

def _audio_to_wav(buf):
    import numpy as np, wave as _wave
    buf_bytes = buf.tobytes() if hasattr(buf, 'tobytes') else buf
    return buf_bytes

# ── Audio rendering helper ──

def render_audio(audio_np):
    import numpy as np, wave as _wave
    buf = io.BytesIO()
    with _wave.open(buf, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(24000)
        pcm = (np.clip(audio_np, -1.0, 1.0) * 32767).astype(np.int16).tobytes()
        wf.writeframes(pcm)
    buf.seek(0)
    return buf

# ── Routes ──

@app.route("/health")
def health():
    ready = engine is not None and engine.ready
    speakers = [s["name"] for s in list_all_speakers()]
    return jsonify({"status": "ok" if ready else "loading",
                    "speakers": speakers, "languages": LANGUAGES})

@app.route("/v1/voices", methods=["GET"])
def voices_list():
    return jsonify({"voices": list_all_speakers()})

@app.route("/v1/voices/<name>", methods=["GET"])
def voices_get(name):
    name = name.lower()
    if name in [s.lower() for s in PRESET_SPEAKERS]:
        return jsonify({"name": name, "type": "preset"})
    if not _voice_exists(name):
        return jsonify({"error": "voice not found"}), 404
    meta = _load_voice_meta(name)
    return jsonify({"name": name, "type": "custom", **meta})

@app.route("/v1/voices/<name>", methods=["DELETE"])
def voices_delete(name):
    name = name.lower()
    if name in [s.lower() for s in PRESET_SPEAKERS]:
        return jsonify({"error": "cannot delete preset voice"}), 400
    path = os.path.join(CUSTOM_VOICES_DIR, f"{name}.json")
    if not os.path.exists(path):
        return jsonify({"error": "voice not found"}), 404
    os.remove(path)
    return jsonify({"deleted": name})

@app.route("/v1/voices/train", methods=["POST"])
def voices_train():
    if engine is None or not engine.ready:
        return jsonify({"error": "engine not ready"}), 503

    if "audio" not in request.files:
        return jsonify({"error": "missing 'audio' file in multipart form"}), 400

    name = request.form.get("name", "").strip().lower()
    if not name or not SAFE_NAME_RE.match(name):
        return jsonify({"error": "invalid name (a-z, 0-9, _, - ; max 64 chars)"}), 400
    if name in [s.lower() for s in PRESET_SPEAKERS]:
        return jsonify({"error": "name conflicts with preset speaker"}), 400

    ref_text = request.form.get("text", "").strip()
    description = request.form.get("description", "").strip()
    language = request.form.get("language", "english").strip().lower()
    if language not in LANGUAGES:
        language = "english"

    audio_file = request.files["audio"]
    allowed_ext = {".wav", ".mp3", ".flac", ".m4a", ".opus", ".ogg"}
    ext = os.path.splitext(audio_file.filename or "upload.wav")[1].lower()
    if ext not in allowed_ext:
        return jsonify({"error": f"unsupported format: {ext}"}), 400

    try:
        import numpy as np
        from qwen3_tts_gguf.inference.utils.audio import load_audio
        from qwen3_tts_gguf.inference import TTSConfig
        from qwen3_tts_gguf.inference.schema.result import TTSResult

        with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp:
            audio_file.save(tmp)
            tmp_path = tmp.name

        samples = load_audio(tmp_path)
        os.unlink(tmp_path)

        if samples is None or len(samples) < 2400:
            return jsonify({"error": "audio too short (min 0.1s at 24kHz)"}), 400

        duration = len(samples) / 24000.0
        if duration > 30.0:
            samples = samples[:int(30.0 * 24000)]
            duration = 30.0

        codes = engine.codec_encoder.encode(samples)
        spk_emb = engine.speaker_encoder.encode(samples)

        text_ids = engine.tokenizer.encode(ref_text).ids if ref_text else []

        result = TTSResult(
            text=ref_text,
            text_ids=text_ids,
            codes=codes,
            spk_emb=spk_emb,
            audio=samples
        )

        voice_data = {
            "name": name,
            "description": description,
            "language": language,
            "created": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "duration_hint": round(duration, 2),
            "text": ref_text,
            "text_ids": text_ids,
            "codes": codes.tolist(),
            "spk_emb": result.spk_emb.tolist(),
        }

        out_path = os.path.join(CUSTOM_VOICES_DIR, f"{name}.json")
        with open(out_path, "w", encoding="utf-8") as f:
            json.dump(voice_data, f, ensure_ascii=False)

        preview_audio = None
        if ref_text:
            stream = engine.create_stream(n_ctx=2048)
            if stream is not None:
                stream.set_voice(result)
                clone_result = stream.clone(text=ref_text, language=language, config=TTSConfig())
                stream.join()
                if clone_result and clone_result.audio is not None:
                    preview_audio = clone_result.audio

        resp = {"name": name, "type": "custom", "description": description,
                "duration": round(duration, 2), "spk_emb_dim": len(spk_emb),
                "codes_frames": len(codes)}

        if preview_audio is not None:
            wav_buf = render_audio(preview_audio)
            resp_json = json.dumps(resp)
            return send_file(wav_buf, mimetype="audio/wav", download_name=f"{name}_preview.wav",
                             as_attachment=False), 200, {"X-Voice-Info": resp_json}

        return jsonify(resp), 201

    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route("/v1/audio/speech", methods=["POST"])
def speech():
    if engine is None or not engine.ready:
        return jsonify({"error": "engine not ready"}), 503
    data = request.get_json(force=True)
    text = data.get("input", "")
    if not text:
        return jsonify({"error": "missing input"}), 400

    voice_name = data.get("voice", "Vivian")
    language = data.get("language", "english")
    if language not in LANGUAGES:
        language = "english"
    instruct = data.get("instruct", "")

    try:
        import numpy as np
        from qwen3_tts_gguf.inference import TTSConfig
        from qwen3_tts_gguf.inference.schema.result import TTSResult

        stream = engine.create_stream(n_ctx=2048)
        if stream is None:
            return jsonify({"error": "failed to create stream"}), 500

        cfg = TTSConfig()
        voice_json = os.path.join(CUSTOM_VOICES_DIR, f"{voice_name.lower()}.json")

        if os.path.exists(voice_json):
            with open(voice_json, "r", encoding="utf-8") as f:
                vdata = json.load(f)
            spk_emb = np.array(vdata["spk_emb"], dtype=np.float32)
            codes = np.array(vdata["codes"], dtype=np.int64)
            anchor = TTSResult(
                text=vdata.get("text", ""),
                text_ids=vdata.get("text_ids", []),
                codes=codes,
                spk_emb=spk_emb
            )
            stream.set_voice(anchor)
            result = stream.clone(text=text, language=language, config=cfg)
        else:
            speaker = voice_name
            if speaker not in PRESET_SPEAKERS:
                speaker = "Vivian"
            result = stream.custom(text=text, speaker=speaker, language=language,
                                   instruct=instruct, config=cfg)

        stream.join()
        if result is None:
            return jsonify({"error": "synthesis returned None"}), 500

        audio = result.audio if hasattr(result, 'audio') and result.audio is not None else None
        if audio is None:
            tmp = os.path.join("/tmp", "tts_out.wav")
            result.save(tmp)
            import soundfile as sf
            audio, _ = sf.read(tmp, dtype='float32')
            os.remove(tmp)

        return send_file(render_audio(audio), mimetype="audio/wav", download_name="speech.wav")
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=8072)
    p.add_argument("--model-dir", default=MODEL_DIR)
    a = p.parse_args()
    MODEL_DIR = a.model_dir
    CUSTOM_VOICES_DIR = os.path.join(MODEL_DIR, "custom_speakers")
    init_engine()
    app.run(host=a.host, port=a.port)