import io
import os
import uuid
import wave
import tempfile
import threading
import urllib.request
from pathlib import Path
from flask import Flask, render_template, request, jsonify, send_file

app = Flask(__name__)

# HuggingFace Spaces stores files under /app
BASE_DIR   = Path(__file__).parent
OUTPUT_DIR = BASE_DIR / "static" / "audio"
MODELS_DIR = BASE_DIR / "models"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
MODELS_DIR.mkdir(parents=True, exist_ok=True)

HF_BASE = "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0"
HF_MAIN = "https://huggingface.co/rhasspy/piper-voices/resolve/main"

# ─────────────────────────────────────────────
# Voice Registry
# ─────────────────────────────────────────────
VOICES = {
    # ── Hindi — on HF_MAIN (added after v1.0.0 tag) ──
    "hi_rohan": {
        "label": "Rohan", "lang": "hi", "gender": "male", "quality": "medium",
        "onnx_url": f"{HF_MAIN}/hi/hi_IN/rohan/medium/hi_IN-rohan-medium.onnx",
        "json_url": f"{HF_MAIN}/hi/hi_IN/rohan/medium/hi_IN-rohan-medium.onnx.json",
        "onnx": "hi_IN-rohan-medium.onnx",
        "json": "hi_IN-rohan-medium.onnx.json",
    },
    "hi_pratham": {
        "label": "Pratham", "lang": "hi", "gender": "male", "quality": "medium",
        "onnx_url": f"{HF_MAIN}/hi/hi_IN/pratham/medium/hi_IN-pratham-medium.onnx",
        "json_url": f"{HF_MAIN}/hi/hi_IN/pratham/medium/hi_IN-pratham-medium.onnx.json",
        "onnx": "hi_IN-pratham-medium.onnx",
        "json": "hi_IN-pratham-medium.onnx.json",
    },
    "hi_priyamvada": {
        "label": "Priyamvada", "lang": "hi", "gender": "female", "quality": "medium",
        "onnx_url": f"{HF_MAIN}/hi/hi_IN/priyamvada/medium/hi_IN-priyamvada-medium.onnx",
        "json_url": f"{HF_MAIN}/hi/hi_IN/priyamvada/medium/hi_IN-priyamvada-medium.onnx.json",
        "onnx": "hi_IN-priyamvada-medium.onnx",
        "json": "hi_IN-priyamvada-medium.onnx.json",
    },
    # ── English US ──
    "en_lessac": {
        "label": "Lessac", "lang": "en", "gender": "male", "quality": "medium",
        "onnx_url": f"{HF_BASE}/en/en_US/lessac/medium/en_US-lessac-medium.onnx",
        "json_url": f"{HF_BASE}/en/en_US/lessac/medium/en_US-lessac-medium.onnx.json",
        "onnx": "en_US-lessac-medium.onnx",
        "json": "en_US-lessac-medium.onnx.json",
    },
    "en_amy": {
        "label": "Amy", "lang": "en", "gender": "female", "quality": "medium",
        "onnx_url": f"{HF_BASE}/en/en_US/amy/medium/en_US-amy-medium.onnx",
        "json_url": f"{HF_BASE}/en/en_US/amy/medium/en_US-amy-medium.onnx.json",
        "onnx": "en_US-amy-medium.onnx",
        "json": "en_US-amy-medium.onnx.json",
    },
    "en_ryan": {
        "label": "Ryan", "lang": "en", "gender": "male", "quality": "high",
        "onnx_url": f"{HF_BASE}/en/en_US/ryan/high/en_US-ryan-high.onnx",
        "json_url": f"{HF_BASE}/en/en_US/ryan/high/en_US-ryan-high.onnx.json",
        "onnx": "en_US-ryan-high.onnx",
        "json": "en_US-ryan-high.onnx.json",
    },
    "en_kristin": {
        "label": "Kristin", "lang": "en", "gender": "female", "quality": "medium",
        "onnx_url": f"{HF_BASE}/en/en_US/kristin/medium/en_US-kristin-medium.onnx",
        "json_url": f"{HF_BASE}/en/en_US/kristin/medium/en_US-kristin-medium.onnx.json",
        "onnx": "en_US-kristin-medium.onnx",
        "json": "en_US-kristin-medium.onnx.json",
    },
    "en_joe": {
        "label": "Joe", "lang": "en", "gender": "male", "quality": "medium",
        "onnx_url": f"{HF_BASE}/en/en_US/joe/medium/en_US-joe-medium.onnx",
        "json_url": f"{HF_BASE}/en/en_US/joe/medium/en_US-joe-medium.onnx.json",
        "onnx": "en_US-joe-medium.onnx",
        "json": "en_US-joe-medium.onnx.json",
    },
    # ── English GB ──
    "en_gb_alan": {
        "label": "Alan (British)", "lang": "en", "gender": "male", "quality": "medium",
        "onnx_url": f"{HF_BASE}/en/en_GB/alan/medium/en_GB-alan-medium.onnx",
        "json_url": f"{HF_BASE}/en/en_GB/alan/medium/en_GB-alan-medium.onnx.json",
        "onnx": "en_GB-alan-medium.onnx",
        "json": "en_GB-alan-medium.onnx.json",
    },
    "en_gb_jenny": {
        "label": "Jenny (British)", "lang": "en", "gender": "female", "quality": "medium",
        "onnx_url": f"{HF_BASE}/en/en_GB/jenny_dioco/medium/en_GB-jenny_dioco-medium.onnx",
        "json_url": f"{HF_BASE}/en/en_GB/jenny_dioco/medium/en_GB-jenny_dioco-medium.onnx.json",
        "onnx": "en_GB-jenny_dioco-medium.onnx",
        "json": "en_GB-jenny_dioco-medium.onnx.json",
    },
}

# ─────────────────────────────────────────────
# State
# ─────────────────────────────────────────────
loaded_voices  = {}
loading_status = {}
tts_lock       = threading.Lock()


# ─────────────────────────────────────────────
# Download & Load
# ─────────────────────────────────────────────
def _download(url, dest):
    if dest.exists():
        return
    print(f"  ⬇  Downloading {dest.name} …", flush=True)
    tmp = dest.with_suffix(".tmp")
    try:
        urllib.request.urlretrieve(url, tmp)
        tmp.rename(dest)
        print(f"  ✅ {dest.name} ready", flush=True)
    except Exception as e:
        tmp.unlink(missing_ok=True)
        raise RuntimeError(f"Download failed: {e}")


def _load_voice(key):
    from piper import PiperVoice
    loading_status[key] = "loading"
    v = VOICES[key]
    onnx_path = MODELS_DIR / v["onnx"]
    json_path = MODELS_DIR / v["json"]
    try:
        _download(v["onnx_url"], onnx_path)
        _download(v["json_url"], json_path)
        voice = PiperVoice.load(str(onnx_path), config_path=str(json_path))
        with tts_lock:
            loaded_voices[key] = voice
        loading_status[key] = "ready"
        print(f"  🎙  Voice '{key}' ready.", flush=True)
    except Exception as e:
        loading_status[key] = f"error: {e}"
        print(f"  ❌ Voice '{key}' failed: {e}", flush=True)


def preload_defaults():
    for key in ("en_lessac", "hi_rohan"):
        t = threading.Thread(target=_load_voice, args=(key,), daemon=True)
        t.start()


# ─────────────────────────────────────────────
# Synthesis — handles all piper-tts API versions
# ─────────────────────────────────────────────
def _synthesize_to_bytes(voice, text):
    # v2 API: synthesize_wav
    if hasattr(voice, "synthesize_wav"):
        buf = io.BytesIO()
        with wave.open(buf, "wb") as wf:
            voice.synthesize_wav(text, wf)
        return buf.getvalue()

    # v1 API: synthesize_stream_raw (yields raw PCM)
    if hasattr(voice, "synthesize_stream_raw"):
        buf = io.BytesIO()
        wf = wave.open(buf, "wb")
        try:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(voice.config.sample_rate)
            for chunk in voice.synthesize_stream_raw(text):
                wf.writeframes(chunk)
        finally:
            wf.close()
        return buf.getvalue()

    # v1 fallback: synthesize(text, wav_file) — needs a real file path
    tmp_fd, tmp_path = tempfile.mkstemp(suffix=".wav")
    try:
        os.close(tmp_fd)
        with wave.open(tmp_path, "w") as wf:
            voice.synthesize(text, wf)
        return Path(tmp_path).read_bytes()
    finally:
        Path(tmp_path).unlink(missing_ok=True)


# ─────────────────────────────────────────────
# Routes
# ─────────────────────────────────────────────
@app.route("/")
def index():
    return render_template("index.html")


@app.route("/status")
def status():
    ready = any(s == "ready" for s in loading_status.values())
    return jsonify({"ready": ready, "loading": loading_status})


@app.route("/voices")
def voices_route():
    return jsonify({"voices": [
        {
            "key":     k,
            "label":   v["label"],
            "lang":    v["lang"],
            "gender":  v["gender"],
            "quality": v["quality"],
            "status":  loading_status.get(k, "not_loaded"),
        }
        for k, v in VOICES.items()
    ]})


@app.route("/load_voice", methods=["POST"])
def load_voice_route():
    key = request.get_json().get("key")
    if key not in VOICES:
        return jsonify({"error": "Unknown voice key"}), 400
    if loading_status.get(key) not in ("ready", "loading"):
        threading.Thread(target=_load_voice, args=(key,), daemon=True).start()
    return jsonify({"status": loading_status.get(key, "loading")})


@app.route("/synthesize", methods=["POST"])
def synthesize():
    data      = request.get_json()
    text      = data.get("text", "").strip()
    voice_key = data.get("voice", "en_lessac")

    if not text:
        return jsonify({"error": "Text cannot be empty"}), 400
    if len(text) > 50000:
        return jsonify({"error": "Text too long (max 50,000 characters)"}), 400
    if voice_key not in VOICES:
        return jsonify({"error": "Unknown voice"}), 400

    if loading_status.get(voice_key) not in ("ready", "loading"):
        threading.Thread(target=_load_voice, args=(voice_key,), daemon=True).start()

    if loading_status.get(voice_key) == "loading":
        return jsonify({"error": "Voice is still loading, please retry in a moment."}), 503
    if loading_status.get(voice_key, "").startswith("error"):
        return jsonify({"error": loading_status[voice_key]}), 500

    voice = loaded_voices.get(voice_key)
    if not voice:
        return jsonify({"error": "Voice not ready yet."}), 503

    file_id  = uuid.uuid4().hex[:12]
    out_path = OUTPUT_DIR / f"speech_{VOICES[voice_key]['lang']}_{file_id}.wav"

    try:
        with tts_lock:
            out_path.write_bytes(_synthesize_to_bytes(voice, text))
        _cleanup_old_files()
        return jsonify({"success": True, "file": f"/audio/{out_path.name}", "filename": out_path.name})
    except Exception as e:
        out_path.unlink(missing_ok=True)
        return jsonify({"error": f"Synthesis failed: {e}"}), 500


@app.route("/audio/<filename>")
def serve_audio(filename):
    path = OUTPUT_DIR / Path(filename).name
    if not path.exists():
        return jsonify({"error": "File not found"}), 404
    return send_file(str(path), mimetype="audio/wav")


@app.route("/download/<filename>")
def download_audio(filename):
    path = OUTPUT_DIR / Path(filename).name
    if not path.exists():
        return jsonify({"error": "File not found"}), 404
    return send_file(str(path), mimetype="audio/wav", as_attachment=True, download_name=filename)


def _cleanup_old_files(keep=20):
    files = sorted(OUTPUT_DIR.glob("speech_*.wav"), key=lambda f: f.stat().st_mtime)
    for f in files[:-keep]:
        f.unlink(missing_ok=True)


if __name__ == "__main__":
    print("\n🎙️  VoxStudio — HuggingFace Space", flush=True)
    print("📦  Pre-loading default voices in background…", flush=True)
    preload_defaults()
    # HuggingFace Spaces requires port 7860
    port = int(os.environ.get("PORT", 7860))
    print(f"🚀  Running on http://0.0.0.0:{port}", flush=True)
    app.run(host="0.0.0.0", port=port, debug=False, threaded=True)