#!/usr/bin/env python3
"""SOTAKA — Speak (VoxCPM2) + Song Convert (Seed-VC) on Hugging Face Spaces."""

from __future__ import annotations

import os

os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")

from pathlib import Path
from typing import Optional, Tuple

import gradio as gr
import numpy as np
import soundfile as sf

try:
    import spaces
except ImportError:  # local / non-ZeroGPU

    class spaces:  # type: ignore
        @staticmethod
        def GPU(*_args, **_kwargs):
            def deco(fn):
                return fn

            return deco


MODEL_ID = os.environ.get("VOXCPM_MODEL", "openbmb/VoxCPM2")
OUT_DIR = Path("outputs")
OUT_DIR.mkdir(exist_ok=True)
# Long text is split into chunks (VoxCPM cache ~8192). Ref still capped.
CHUNK_CHARS = int(os.environ.get("SOTAKA_CHUNK_CHARS", "280"))
MAX_SPEAK_TEXT_CHARS = int(os.environ.get("SOTAKA_MAX_SPEAK_CHARS", "12000"))
MAX_REF_SEC = float(os.environ.get("SOTAKA_MAX_REF_SEC", "12"))
# Chunks per ZeroGPU session (~120s wall clock; ~5 × ~15–20s gen)
CHUNKS_PER_GPU = int(os.environ.get("SOTAKA_CHUNKS_PER_GPU", "5"))

_model = None

SOTAKA_CSS = """
/* Progress fill Gradio uses: .eta-bar { background: var(--background-fill-secondary) } */
/* Force orange on that bar (the gray strip that grows left→right while processing) */

.generating,
.wrap.generating,
div[class*="generating"] {
  --background-fill-secondary: #ea580c !important;
  --loader-color: #ea580c !important;
  --color-accent: #ea580c !important;
}

.eta-bar,
div.eta-bar,
[class*="eta-bar"] {
  background: #ea580c !important;
  background-color: #ea580c !important;
  background-image: linear-gradient(
    90deg,
    #fdba74 0%,
    #ea580c 55%,
    #c2410c 100%
  ) !important;
  opacity: 1 !important;
  z-index: 6 !important;
}

.progress-bar,
[class*="progress-bar"]:not([class*="wrap"]) {
  background: #ea580c !important;
  background-color: #ea580c !important;
}

.progress-bar-wrap,
[class*="progress-bar-wrap"] {
  border: 2px solid #ea580c !important;
  background: rgba(234, 88, 12, 0.15) !important;
}

.meta-text,
[class*="meta-text"] {
  color: #9a3412 !important;
  background: rgba(255, 247, 237, 0.95) !important;
  font-weight: 700 !important;
}

/* Phone */
@media (max-width: 768px) {
  #sotaka-app .row,
  #sotaka-app .form,
  .gradio-container .row,
  .gradio-container .form {
    flex-direction: column !important;
    gap: 10px !important;
  }
  #sotaka-app .row > *,
  .gradio-container .row > * {
    width: 100% !important;
    max-width: 100% !important;
    flex: 1 1 100% !important;
    min-width: 0 !important;
  }
  .gradio-container button.primary,
  #sotaka-app button {
    width: 100% !important;
    min-height: 48px !important;
    font-size: 16px !important;
  }
  .gradio-container textarea,
  .gradio-container input {
    font-size: 16px !important;
  }
  .gradio-container audio {
    max-width: 100% !important;
  }
}
"""


def get_model():
    global _model
    if _model is None:
        from voxcpm import VoxCPM

        _model = VoxCPM.from_pretrained(MODEL_ID, load_denoiser=False)
    return _model


def _build_text(text: str, instruct: Optional[str]) -> str:
    if instruct:
        return f"({instruct}){text}"
    return text


def _split_long_text(text: str, max_chars: int = CHUNK_CHARS) -> list:
    """Split long text on sentence boundaries for stable VoxCPM generation."""
    import re

    text = re.sub(r"\s+", " ", text).strip()
    if len(text) <= max_chars:
        return [text]

    # Khmer ។ + Latin . ! ? and newlines
    pieces = re.split(r"(?<=[។!?\.…])\s+|\n+", text)
    chunks: list[str] = []
    buf = ""
    for p in pieces:
        p = p.strip()
        if not p:
            continue
        if len(p) > max_chars:
            if buf:
                chunks.append(buf.strip())
                buf = ""
            for i in range(0, len(p), max_chars):
                chunks.append(p[i : i + max_chars].strip())
            continue
        if not buf:
            buf = p
        elif len(buf) + 1 + len(p) <= max_chars:
            buf = f"{buf} {p}"
        else:
            chunks.append(buf.strip())
            buf = p
    if buf.strip():
        chunks.append(buf.strip())
    return [c for c in chunks if c]


def _trim_ref_wav(path: str, max_sec: float = MAX_REF_SEC) -> Tuple[str, str]:
    """Trim reference clip; long refs cause tensor size 8192 vs 16k+ crashes."""
    import librosa

    y, sr = librosa.load(path, sr=None, mono=True)
    dur = len(y) / float(sr)
    note = ""
    if dur > max_sec:
        y = y[: int(sr * max_sec)]
        out = OUT_DIR / "speak_ref_trim.wav"
        sf.write(out, y, sr)
        note = f"Reference trimmed {dur:.1f}s → {max_sec:.0f}s.\n"
        return str(out), note
    return path, note


def _toast(path: Optional[str], log: str, ok: str) -> Tuple[Optional[str], str]:
    """Screen alert for success / failure (Speak + Song)."""
    brief = (log or "").strip().split("\n")[0][:220] or "Unknown error"
    if path:
        gr.Success(ok)
        if log:
            gr.Info(brief)
    else:
        gr.Warning(brief)
    return path, log


def _clone_kwargs(ref_audio: Optional[str], ref_text: Optional[str]) -> dict:
    if not ref_audio:
        return {}
    if ref_text:
        return {"prompt_wav_path": ref_audio, "prompt_text": ref_text}
    return {"reference_wav_path": ref_audio}


# ZeroGPU: only ONE @spaces.GPU call per user click (token expires on 2nd call).
_CHUNK_SEP = "\n\x1e\n"


@spaces.GPU(duration=120)
def _speak_batch(
    chunks_joined: str,
    instruct: str,
    ref_audio: str,
    ref_text: str,
    timesteps: int,
    cfg: float,
    batch_index: int,
) -> Tuple[Optional[str], str]:
    """One ZeroGPU session: generate a batch, save wav, return filepath."""
    global _model
    _model = None
    chunks = [c.strip() for c in chunks_joined.split(_CHUNK_SEP) if c.strip()]
    if not chunks:
        return None, "Empty batch."

    try:
        model = get_model()
    except Exception as e:
        return None, f"Failed to load model: {e}"

    instruct_opt = (instruct or "").strip() or None
    ref_opt = (ref_audio or "").strip() or None
    ref_text_opt = (ref_text or "").strip() or None
    clone = _clone_kwargs(ref_opt, ref_text_opt)

    parts: list[np.ndarray] = []
    sample_rate = 48000
    for i, chunk in enumerate(chunks):
        kwargs = {
            "text": _build_text(chunk, instruct_opt),
            "inference_timesteps": int(timesteps),
            "cfg_value": float(cfg),
            "max_len": 768,
        }
        kwargs.update(clone)
        try:
            wav = model.generate(**kwargs)
            if wav is None or (hasattr(wav, "size") and wav.size == 0):
                return None, f"No audio for batch chunk {i+1}/{len(chunks)}"
            wav = np.asarray(wav, dtype=np.float32).reshape(-1)
            sample_rate = int(getattr(model.tts_model, "sample_rate", 48000))
            parts.append(wav)
            if i < len(chunks) - 1:
                parts.append(np.zeros(int(sample_rate * 0.25), dtype=np.float32))
        except Exception as e:
            msg = str(e)
            if "expanded size of the tensor" in msg or "8192" in msg:
                msg = f"Chunk too long for model cache. Detail: {e}"
            if "Expired ZeroGPU" in msg or "proxy token" in msg.lower():
                msg = (
                    "ZeroGPU session expired. Click Generate / Continue once "
                    f"(one GPU run per click). Detail: {e}"
                )
            return None, msg

    out = OUT_DIR / f"speak_batch_{int(batch_index)}.wav"
    sf.write(out, np.concatenate(parts), sample_rate)
    return str(out), f"ok sr={sample_rate} n={len(chunks)}"


def _merge_wavs(paths: list, gap_sec: float = 0.2) -> Tuple[str, float, int]:
    import librosa

    parts: list[np.ndarray] = []
    sr = 48000
    for i, p in enumerate(paths):
        y, sr = librosa.load(p, sr=None, mono=True)
        parts.append(np.asarray(y, dtype=np.float32))
        if i < len(paths) - 1:
            parts.append(np.zeros(int(sr * gap_sec), dtype=np.float32))
    wav = np.concatenate(parts)
    out = OUT_DIR / "SOTAKA_Voice.wav"
    sf.write(out, wav, int(sr))
    return str(out), len(wav) / float(sr), int(sr)


def synthesize(
    text: str,
    instruct: str,
    ref_audio: Optional[str],
    ref_text: str,
    timesteps: int,
    cfg: float,
):
    """First GPU run only — use Continue for remaining chunks (fresh ZeroGPU token)."""
    text = (text or "").strip()
    if not text:
        empty = {"remaining": [], "parts": [], "meta": {}}
        return _toast(None, "Please enter text.", "") + (
            empty,
            gr.update(interactive=False),
        )

    notes = []
    if len(text) > MAX_SPEAK_TEXT_CHARS:
        text = text[:MAX_SPEAK_TEXT_CHARS]
        notes.append(f"Text capped at {MAX_SPEAK_TEXT_CHARS} chars (safety).")

    instruct = (instruct or "").strip()
    ref_text = (ref_text or "").strip()
    ref_path = ""
    if ref_audio:
        try:
            ref_path, trim_note = _trim_ref_wav(ref_audio, MAX_REF_SEC)
            if trim_note:
                notes.append(trim_note.strip())
        except Exception as e:
            empty = {"remaining": [], "parts": [], "meta": {}}
            return _toast(None, f"Could not read reference audio: {e}", "") + (
                empty,
                gr.update(interactive=False),
            )

    chunks = _split_long_text(text, CHUNK_CHARS)
    batch = chunks[:CHUNKS_PER_GPU]
    remaining = chunks[CHUNKS_PER_GPU:]
    notes.append(
        f"Long text → {len(chunks)} chunks. This run: {len(batch)}. "
        f"Remaining after: {len(remaining)} (click Continue)."
        if remaining
        else f"Text → {len(chunks)} chunk(s) in this run."
    )

    try:
        path, err = _speak_batch(
            _CHUNK_SEP.join(batch),
            instruct,
            ref_path or "",
            ref_text,
            int(timesteps),
            float(cfg),
            0,
        )
    except Exception as e:
        empty = {"remaining": [], "parts": [], "meta": {}}
        return _toast(None, f"Generation error: {e}", "") + (
            empty,
            gr.update(interactive=False),
        )

    if not path:
        empty = {"remaining": [], "parts": [], "meta": {}}
        return _toast(None, f"Generation error: {err}", "") + (
            empty,
            gr.update(interactive=False),
        )

    parts = [path]
    out, dur, sr = _merge_wavs(parts)
    state = {
        "remaining": remaining,
        "parts": parts,
        "meta": {
            "instruct": instruct,
            "ref_path": ref_path or "",
            "ref_text": ref_text,
            "timesteps": int(timesteps),
            "cfg": float(cfg),
            "total": len(chunks),
            "done": len(batch),
            "batch_i": 1,
        },
    }
    log = "\n".join(notes) + f"\n{err}\n"
    log += f"Model: {MODEL_ID}\nchunks_done={len(batch)}/{len(chunks)} | duration={dur:.2f}s | sr={sr}"
    if remaining:
        gr.Info(f"Partial OK — click Continue for {len(remaining)} more chunks.")
        ok = f"Partial Speak — {dur:.1f}s ({len(batch)}/{len(chunks)}). Click Continue."
    else:
        ok = f"Speak done — {dur:.1f}s ({len(chunks)} chunk(s))"
    audio, log_out = _toast(out, log, ok)
    return audio, log_out, state, gr.update(interactive=bool(remaining))


def continue_speak(state: Optional[dict]):
    """Next ZeroGPU run with a fresh proxy token (one @spaces.GPU call)."""
    empty = {"remaining": [], "parts": [], "meta": {}}
    if not state or not state.get("remaining"):
        return _toast(None, "Nothing left to Continue. Click Generate first.", "") + (
            state or empty,
            gr.update(interactive=False),
        )

    remaining = list(state["remaining"])
    meta = state.get("meta") or {}
    batch = remaining[:CHUNKS_PER_GPU]
    left = remaining[CHUNKS_PER_GPU:]
    batch_i = int(meta.get("batch_i", 1))

    try:
        path, err = _speak_batch(
            _CHUNK_SEP.join(batch),
            meta.get("instruct", ""),
            meta.get("ref_path", ""),
            meta.get("ref_text", ""),
            int(meta.get("timesteps", 5)),
            float(meta.get("cfg", 2.0)),
            batch_i,
        )
    except Exception as e:
        return _toast(None, f"Generation error: {e}", "") + (
            state,
            gr.update(interactive=True),
        )

    if not path:
        return _toast(None, f"Generation error: {err}", "") + (
            state,
            gr.update(interactive=True),
        )

    parts = list(state.get("parts") or []) + [path]
    out, dur, sr = _merge_wavs(parts)
    done = int(meta.get("done", 0)) + len(batch)
    total = int(meta.get("total", done))
    new_state = {
        "remaining": left,
        "parts": parts,
        "meta": {**meta, "done": done, "batch_i": batch_i + 1},
    }
    log = (
        f"Continue batch OK ({len(batch)} chunks). {err}\n"
        f"Model: {MODEL_ID}\n"
        f"chunks_done={done}/{total} | duration={dur:.2f}s | sr={sr}"
    )
    if left:
        gr.Info(f"Still {len(left)} chunks left — click Continue again.")
        ok = f"Partial Speak — {dur:.1f}s ({done}/{total}). Click Continue."
    else:
        ok = f"Speak done — {dur:.1f}s ({total} chunks)"
    audio, log_out = _toast(out, log, ok)
    return audio, log_out, new_state, gr.update(interactive=bool(left))


from song_convert import convert_song as _convert_song  # noqa: E402
from vocal_remove import separate_vocals as _separate_vocals  # noqa: E402


def convert_song_ui(
    source_audio: Optional[str],
    voice_ref: Optional[str],
    diffusion_steps: int,
    pitch_shift: int,
    length_adjust: float,
    cfg_rate: float,
    auto_f0_adjust: bool = False,
) -> Tuple[Optional[str], str]:
    path, log = _convert_song(
        source_audio,
        voice_ref,
        diffusion_steps,
        pitch_shift,
        length_adjust,
        cfg_rate,
        auto_f0_adjust,
    )
    if path:
        return _toast(path, log, "Song Convert done — download SOTAKA_Song.wav")
    return _toast(None, log or "Song Convert failed", "")


def separate_ui(audio_path: Optional[str], model_name: str):
    try:
        vocals, instrumental, log = _separate_vocals(audio_path, model_name)
    except Exception as e:
        msg = str(e)
        low = msg.lower()
        if "zerogpu" in low or "expired" in low or "gpu" in low:
            msg = (
                "GPU timed out on this song. Keep the tab open and click "
                f"Separate again (long songs are split into chunks). Detail: {e}"
            )
        else:
            msg = f"Separation error: {e}"
        gr.Warning(msg[:220])
        return None, None, msg
    if vocals:
        gr.Success("Vocal Remover done — Vocals + Instrumental ready")
        if log:
            gr.Info(log.strip().split("\n")[0][:220])
        return vocals, instrumental, log
    gr.Warning((log or "Separate failed").split("\n")[0][:220])
    return None, None, log


def build_ui() -> gr.Blocks:
    with gr.Blocks(
        title="SOTAKA",
        css=SOTAKA_CSS,
        elem_id="sotaka-app",
    ) as demo:
        gr.Markdown(
            """
# SOTAKA
**Speak** · **Vocal Remover** · **Song Convert**
            """
        )
        with gr.Tabs():
            with gr.Tab("Speak"):
                gr.Markdown(
                    """
Type Khmer (or other) text — **long text OK**.  
Each **Generate / Continue** = one ZeroGPU run (~5 chunks). Click **Continue** until done.  
Reference clone: **≤12s** clear speech (not a full song).
                    """
                )
                speak_state = gr.State(
                    {"remaining": [], "parts": [], "meta": {}}
                )
                text = gr.Textbox(
                    label="Text (long text OK — use Continue for more)",
                    lines=8,
                    value="សួស្តី។ ខ្ញុំសប្បាយចិត្តដែលបានជួបអ្នក។",
                )
                with gr.Row():
                    instruct = gr.Textbox(
                        label="Voice design (optional)",
                        placeholder="A calm adult Khmer male voice, clear and natural",
                        lines=2,
                    )
                    ref_audio = gr.Audio(
                        label="Reference audio (optional clone)",
                        type="filepath",
                    )
                ref_text = gr.Textbox(
                    label="Reference transcript (optional)",
                    lines=2,
                )
                with gr.Row():
                    timesteps = gr.Slider(4, 16, value=5, step=1, label="Inference steps")
                    cfg = gr.Slider(1.0, 4.0, value=2.0, step=0.1, label="CFG")
                with gr.Row():
                    btn_speak = gr.Button("Generate", variant="primary")
                    btn_continue = gr.Button(
                        "Continue remaining text",
                        variant="secondary",
                        interactive=False,
                    )
                speak_out = gr.Audio(
                    label="Output",
                    type="filepath",
                )
                speak_log = gr.Textbox(
                    label="Log",
                    lines=5,
                )
                btn_speak.click(
                    synthesize,
                    inputs=[text, instruct, ref_audio, ref_text, timesteps, cfg],
                    outputs=[speak_out, speak_log, speak_state, btn_continue],
                )
                btn_continue.click(
                    continue_speak,
                    inputs=[speak_state],
                    outputs=[speak_out, speak_log, speak_state, btn_continue],
                )

            with gr.Tab("Vocal Remover"):
                gr.Markdown(
                    """
Split a mix into **Vocals** + **Instrumental** (Demucs).  
Use **Vocals** as Song Convert source. Songs up to **~5 min** are split into chunks — keep this tab open (a 4 min mix can take several minutes).
                    """
                )
                vr_in = gr.Audio(label="Song / mix", type="filepath")
                vr_model = gr.Dropdown(
                    choices=["htdemucs", "htdemucs_ft"],
                    value="htdemucs",
                    label="Model",
                )
                btn_vr = gr.Button("Separate vocals / music", variant="primary")
                with gr.Row():
                    vr_vocals = gr.Audio(label="Vocals", type="filepath")
                    vr_inst = gr.Audio(label="Instrumental", type="filepath")
                vr_log = gr.Textbox(label="Log", lines=4)
                btn_vr.click(
                    separate_ui,
                    inputs=[vr_in, vr_model],
                    outputs=[vr_vocals, vr_inst, vr_log],
                )

            with gr.Tab("Song Convert"):
                gr.Markdown(
                    """
Keep **melody & timing**, swap to **your voice** (Seed-VC latest singing model).

**For clear voice:**
1. Source = **Vocals only** (run **Vocal Remover** first — never full mix with drums)  
2. Reference = **8–20s** of *your* clear speech/singing, little noise  
3. Diffusion steps **50–80** (default 50; higher = clearer, slower)  
4. Keep this tab open — long songs are split into ~30s GPU chunks
                    """
                )
                with gr.Row():
                    source = gr.Audio(
                        label="Source vocal (use Vocal Remover first)",
                        type="filepath",
                    )
                    voice_ref = gr.Audio(
                        label="Your voice reference (8–20s clear)",
                        type="filepath",
                    )
                with gr.Row():
                    steps = gr.Slider(
                        10,
                        100,
                        value=50,
                        step=1,
                        label="Diffusion steps (50–80 = clearer)",
                    )
                    pitch = gr.Slider(
                        -12, 12, value=0, step=1, label="Pitch shift (semitones)"
                    )
                with gr.Row():
                    length = gr.Slider(0.5, 2.0, value=1.0, step=0.1, label="Length adjust")
                    song_cfg = gr.Slider(0.0, 1.0, value=0.7, step=0.1, label="CFG rate")
                auto_f0 = gr.Checkbox(
                    label="Auto F0 adjust (off for singing; on if pitch range differs a lot)",
                    value=False,
                )
                btn_song = gr.Button("Convert to my voice", variant="primary")
                song_out = gr.Audio(
                    label="Output",
                    type="filepath",
                )
                song_log = gr.Textbox(
                    label="Log",
                    lines=5,
                )
                btn_song.click(
                    convert_song_ui,
                    inputs=[source, voice_ref, steps, pitch, length, song_cfg, auto_f0],
                    outputs=[song_out, song_log],
                )
    return demo


demo = build_ui()
demo.queue(max_size=8)

if __name__ == "__main__":
    print(f"Starting UI. Speak model: {MODEL_ID}")
    demo.launch(ssr_mode=False)