"""
Tool ASR Generate - Batch Text-to-Speech with VoxCPM2 (ZeroGPU Space)
Focused on Khmer (km) and English (en).

Upload a CSV / JSON / XLSX file containing a "text" column, and this app will
generate speech for every row using openbmb/VoxCPM2, then package all the
generated audio + a metadata.csv/json manifest into a downloadable ZIP that is
already structured the way Hugging Face `datasets` expects an audio folder
(audio/<file>.wav + metadata.csv with a "file_name" column), so you can push
it straight to a Dataset repo.

This version runs on a ZeroGPU Space. The model is instantiated at import
time (on CPU/meta); actual CUDA execution only happens inside the function
decorated with @spaces.GPU, which is where the GPU is allocated per call.

Language handling:
    - VoxCPM2 auto-detects language from the text itself (no language tag
      needed). Each row is still labeled "km" or "en" in the output metadata
      (either from an optional "language" column, or auto-detected from the
      Khmer Unicode script range), which is useful for filtering the dataset
      later, but generation itself uses a single shared reference voice.
    - The sidebar lets you upload one optional reference voice clip (+
      transcript) used for cloning across every row, regardless of language.

Input file requirements:
    - text                (required)  the sentence/paragraph to synthesize
    - id                  (optional)  used as the output filename, auto-generated if missing
    - language            (optional)  "en" or "km" — auto-detected if omitted
    - voice_description   (optional)  natural language voice style, e.g. "a calm young woman"
                                       (will be prepended as "(voice_description)text",
                                       overrides the reference-audio cloning for that row)
"""

import os
import io
import json
import uuid
import zipfile
import tempfile

import numpy as np
import pandas as pd
import soundfile as sf
import gradio as gr
import spaces
import torch

from voxcpm import VoxCPM

MODEL_ID = "openbmb/VoxCPM2"

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------
# On ZeroGPU Spaces the model is instantiated at import time (on CPU/meta),
# and actual CUDA execution only happens inside functions decorated with
# @spaces.GPU. Do NOT move the model to cuda manually here.
print(f"Loading {MODEL_ID} ...")
model = VoxCPM.from_pretrained(MODEL_ID, load_denoiser=False)
SAMPLE_RATE = getattr(model.tts_model, "sample_rate", 48000)
print(f"Model loaded. Sample rate = {SAMPLE_RATE}")


# ---------------------------------------------------------------------------
# Language detection (Khmer vs English)
# ---------------------------------------------------------------------------
KHMER_RANGE = (0x1780, 0x17FF)  # Khmer Unicode block


def detect_language(text: str) -> str:
    for ch in str(text):
        if KHMER_RANGE[0] <= ord(ch) <= KHMER_RANGE[1]:
            return "km"
    return "en"


def normalize_language(value: str) -> str:
    v = str(value).strip().lower()
    if v in ("km", "kh", "khm", "khmer", "cambodian"):
        return "km"
    if v in ("en", "eng", "english"):
        return "en"
    return ""  # unrecognized -> fall back to auto-detect


# ---------------------------------------------------------------------------
# File parsing helpers
# ---------------------------------------------------------------------------
def load_table(file_path: str) -> pd.DataFrame:
    ext = os.path.splitext(file_path)[1].lower()

    if ext == ".csv":
        df = pd.read_csv(file_path)
    elif ext in (".xlsx", ".xls"):
        df = pd.read_excel(file_path)
    elif ext == ".json":
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)
        if isinstance(data, dict):
            for key in ("items", "data", "rows"):
                if key in data and isinstance(data[key], list):
                    data = data[key]
                    break
        df = pd.DataFrame(data)
    else:
        raise ValueError(f"Unsupported file type: {ext}. Use .csv, .json, or .xlsx")

    df.columns = [str(c).strip().lower() for c in df.columns]

    if "text" not in df.columns:
        for alt in ("content", "sentence", "script", "transcript"):
            if alt in df.columns:
                df = df.rename(columns={alt: "text"})
                break

    if "text" not in df.columns:
        raise ValueError("No 'text' column found in the uploaded file.")

    if "id" not in df.columns:
        df.insert(0, "id", [f"{i + 1:04d}" for i in range(len(df))])
    else:
        df["id"] = df["id"].astype(str)

    if "voice_description" not in df.columns:
        df["voice_description"] = ""
    df["voice_description"] = df["voice_description"].fillna("").astype(str)

    df["text"] = df["text"].astype(str)
    df = df[df["text"].str.strip() != ""].reset_index(drop=True)

    if "language" in df.columns:
        df["language"] = df["language"].fillna("").apply(normalize_language)
    else:
        df["language"] = ""
    # fill anything unrecognized/empty with auto-detection from the text script
    needs_detect = df["language"] == ""
    df.loc[needs_detect, "language"] = df.loc[needs_detect, "text"].apply(detect_language)

    return df


def preview_file(file):
    if file is None:
        return None, "Upload a CSV, JSON, or XLSX file with a `text` column."
    try:
        df = load_table(file)
    except Exception as e:
        return None, f"❌ Error reading file: {e}"
    preview = df[["id", "language", "text", "voice_description"]].head(20)
    n_en = int((df["language"] == "en").sum())
    n_km = int((df["language"] == "km").sum())
    return preview, f"✅ Loaded **{len(df)}** rows — English: {n_en}, Khmer: {n_km}. Showing first {min(20, len(df))}."


# ---------------------------------------------------------------------------
# GPU-bound single-item generation
# ---------------------------------------------------------------------------
@spaces.GPU(duration=120)
def _generate_one(text, voice_description, ref_wav_path, ref_text, cfg_value, inference_timesteps):
    prompt = text
    if voice_description and voice_description.strip():
        prompt = f"({voice_description.strip()}){text}"

    kwargs = dict(
        text=prompt,
        cfg_value=float(cfg_value),
        inference_timesteps=int(inference_timesteps),
    )

    # voice_description (Voice Design) and reference-audio cloning are
    # alternative modes — only clone when no per-row description was given.
    if ref_wav_path and not (voice_description and voice_description.strip()):
        kwargs["reference_wav_path"] = ref_wav_path
        if ref_text and ref_text.strip():
            # transcript + reference wav given -> ultimate cloning mode
            kwargs["prompt_wav_path"] = ref_wav_path
            kwargs["prompt_text"] = ref_text.strip()

    wav = model.generate(**kwargs)
    return np.asarray(wav)


# ---------------------------------------------------------------------------
# Batch orchestration (runs on CPU, calls the GPU function per row so we are
# not bound by a single ZeroGPU call's duration limit)
# ---------------------------------------------------------------------------
N_SAMPLES = 10  # how many generated rows to preview inline after a run


def run_batch(
    file,
    ref_audio,
    ref_text,
    cfg_value,
    inference_timesteps,
    progress=gr.Progress(),
):
    if file is None:
        raise gr.Error("Please upload a text file (CSV / JSON / XLSX) first.")

    df = load_table(file)
    if len(df) == 0:
        raise gr.Error("No valid text rows found in the uploaded file.")

    work_dir = tempfile.mkdtemp(prefix="voxcpm_batch_")
    audio_dir = os.path.join(work_dir, "audio")
    os.makedirs(audio_dir, exist_ok=True)

    records = []
    errors = []
    sample_pairs = []  # [(text, wav_path), ...] for the first N_SAMPLES successes

    rows = list(df.iterrows())
    for _, row in progress.tqdm(rows, desc="Generating speech"):
        rid = str(row["id"])
        text = row["text"]
        lang = row["language"]  # "en" or "km", kept for metadata/labeling only
        vdesc = row.get("voice_description", "")
        fname = f"{rid}.wav"
        fpath = os.path.join(audio_dir, fname)

        try:
            wav = _generate_one(text, vdesc, ref_audio, ref_text, cfg_value, inference_timesteps)
            sf.write(fpath, wav, SAMPLE_RATE)
            duration = float(len(wav)) / SAMPLE_RATE
            records.append(
                {
                    "file_name": f"audio/{fname}",
                    "id": rid,
                    "language": lang,
                    "text": text,
                    "voice_description": vdesc,
                    "duration_sec": round(duration, 3),
                }
            )
            if len(sample_pairs) < N_SAMPLES:
                sample_pairs.append((text, fpath))
        except Exception as e:
            errors.append(f"{rid} ({lang}): {e}")

    meta_df = pd.DataFrame(records)
    meta_csv_path = os.path.join(work_dir, "metadata.csv")
    meta_df.to_csv(meta_csv_path, index=False, encoding="utf-8")

    meta_json_path = os.path.join(work_dir, "metadata.json")
    meta_df.to_json(meta_json_path, orient="records", force_ascii=False, indent=2)

    zip_path = os.path.join(tempfile.gettempdir(), f"voxcpm_output_{uuid.uuid4().hex[:8]}.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for root, _, files in os.walk(work_dir):
            for fn in files:
                full = os.path.join(root, fn)
                arc = os.path.relpath(full, work_dir)
                zf.write(full, arc)

    status = f"✅ Done: **{len(records)}** file(s) generated."
    if errors:
        status += f"\n\n⚠️ {len(errors)} failed:\n" + "\n".join(f"- {e}" for e in errors[:15])

    # build fixed-length update lists for the N_SAMPLES (text, audio) preview slots
    sample_updates = []
    for i in range(N_SAMPLES):
        if i < len(sample_pairs):
            t, wav_path = sample_pairs[i]
            sample_updates.append(gr.update(value=t, visible=True))
            sample_updates.append(gr.update(value=wav_path, visible=True))
        else:
            sample_updates.append(gr.update(value="", visible=False))
            sample_updates.append(gr.update(value=None, visible=False))

    return [zip_path, meta_df, status] + sample_updates


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Tool ASR Generate") as demo:

    with gr.Sidebar():
        gr.Markdown("## 🎙️ Reference Voice")
        gr.Markdown(
            "Optional. Upload a short clip to clone that voice for every row "
            "(works for both English and Khmer text). A per-row "
            "`voice_description` in your file overrides this for that row."
        )
        ref_audio = gr.Audio(label="Reference clip", type="filepath")
        ref_text = gr.Textbox(
            label="Reference transcript (optional, improves cloning quality)",
            placeholder="Exact transcript of the reference clip above",
            lines=3,
        )
        gr.Markdown("## ⚙️ Generation settings")
        cfg_value = gr.Slider(0.5, 4.0, value=2.0, step=0.1, label="CFG value")
        inference_timesteps = gr.Slider(4, 30, value=10, step=1, label="Inference timesteps")

    gr.Markdown(
        """
        # 🌍 Tool ASR Generate — Batch Text-to-Speech (VoxCPM2, ZeroGPU)
        ### Focused on 🇰🇭 Khmer and 🇬🇧 English
        Upload a **CSV / JSON / XLSX** file with a `text` column and generate speech for
        every row using [openbmb/VoxCPM2](https://huggingface.co/openbmb/VoxCPM2).
        Set the reference voice and generation settings in the sidebar. When it's done,
        download a ZIP with all the audio + a `metadata.csv` / `metadata.json` manifest,
        ready to push to a Hugging Face Dataset repo.
        """
    )

    file_in = gr.File(
        label="Upload text file (.csv / .json / .xlsx)",
        file_types=[".csv", ".json", ".xlsx", ".xls"],
    )

    preview_df = gr.Dataframe(label="Preview (first 20 rows)", interactive=False)
    preview_status = gr.Markdown()
    file_in.change(preview_file, inputs=file_in, outputs=[preview_df, preview_status])

    generate_btn = gr.Button("🚀 Generate All", variant="primary")

    result_zip = gr.File(label="Download ZIP (audio/ + metadata.csv + metadata.json)")
    result_meta = gr.Dataframe(label="Generated metadata")
    run_status = gr.Markdown()

    gr.Markdown("### 🔊 Sample preview (first 10 generated rows)")
    sample_components = []
    for i in range(N_SAMPLES):
        with gr.Row():
            t = gr.Textbox(label=f"Text {i + 1}", interactive=False, scale=3, visible=False)
            a = gr.Audio(label=f"Audio {i + 1}", interactive=False, scale=2, visible=False)
        sample_components.append(t)
        sample_components.append(a)

    generate_btn.click(
        run_batch,
        inputs=[file_in, ref_audio, ref_text, cfg_value, inference_timesteps],
        outputs=[result_zip, result_meta, run_status] + sample_components,
    )

demo.queue(max_size=10, default_concurrency_limit=1).launch()