"""
Tool ASR Generate - Batch Text-to-Speech with OmniVoice (CPU 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 k2-fsa/OmniVoice, 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 entirely on CPU (no ZeroGPU / spaces.GPU dependency), so it
works on a plain CPU Space. OmniVoice is a diffusion-language-model-style TTS
model; CPU inference is slower than GPU, though its architecture is
comparatively fast (RTF as low as 0.025 on GPU) — actual CPU throughput will
depend on the Space's hardware. For large batches, consider running fewer,
shorter texts per run.

Language handling:
    - OmniVoice supports 600+ languages and 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/instruction.
    - 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-design attributes, e.g.
                                       "female, low pitch, british accent"
                                       (passed to OmniVoice's `instruct` voice-design mode,
                                       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 torch

from omnivoice import OmniVoice

MODEL_ID = "k2-fsa/OmniVoice"
SAMPLE_RATE = 24000  # OmniVoice always generates 24 kHz audio

# ---------------------------------------------------------------------------
# Model loading (CPU only)
# ---------------------------------------------------------------------------
# Force CPU execution — no CUDA / ZeroGPU involved.
torch.set_num_threads(os.cpu_count() or 4)

print(f"Loading {MODEL_ID} on CPU ...")
model = OmniVoice.from_pretrained(MODEL_ID, device_map="cpu", dtype=torch.float32)
print(f"Model loaded on CPU. 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))}."


# ---------------------------------------------------------------------------
# Single-item generation (runs on CPU)
# ---------------------------------------------------------------------------
def _generate_one(text, voice_description, ref_wav_path, ref_text, num_step, speed):
    kwargs = dict(
        text=text,
        num_step=int(num_step),
        speed=float(speed),
    )

    # voice_description (Voice Design via `instruct`) and reference-audio
    # cloning are alternative modes — only clone when no per-row description
    # was given. If neither is given, OmniVoice falls back to "auto voice".
    if voice_description and voice_description.strip():
        kwargs["instruct"] = voice_description.strip()
    elif ref_wav_path:
        kwargs["ref_audio"] = ref_wav_path
        if ref_text and ref_text.strip():
            kwargs["ref_text"] = ref_text.strip()
        # else: OmniVoice auto-transcribes the reference clip with Whisper.

    with torch.no_grad():
        audio = model.generate(**kwargs)
    # audio is a list of np.ndarray (one per input); we generate one at a time.
    return np.asarray(audio[0])


# ---------------------------------------------------------------------------
# Batch orchestration (sequential, CPU)
# ---------------------------------------------------------------------------
N_SAMPLES = 10  # how many generated rows to preview inline after a run


def run_batch(
    file,
    ref_audio,
    ref_text,
    num_step,
    speed,
    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, num_step, speed)
            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 3–10s clip to clone that voice for every row "
            "(works for both English and Khmer text). A per-row "
            "`voice_description` in your file (e.g. `female, low pitch, british accent`) "
            "switches that row to OmniVoice's Voice Design mode instead."
        )
        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")
        num_step = gr.Slider(4, 64, value=32, step=1, label="Diffusion steps (num_step)")
        speed = gr.Slider(0.5, 2.0, value=1.0, step=0.05, label="Speed")

    gr.Markdown(
        """
        # 🌍 Tool ASR Generate — Batch Text-to-Speech (OmniVoice, CPU)
        ### Focused on 🇰🇭 Khmer and 🇬🇧 English
        Upload a **CSV / JSON / XLSX** file with a `text` column and generate speech for
        every row using [k2-fsa/OmniVoice](https://huggingface.co/k2-fsa/OmniVoice).
        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, num_step, speed],
        outputs=[result_zip, result_meta, run_status] + sample_components,
    )

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