import os
import re
import tempfile
from typing import Optional, Tuple

import gradio as gr
import torch
import torchaudio as ta
from chatterbox.tts_turbo import ChatterboxTurboTTS


MAX_CHARS = 2000
APP_DIR = os.path.dirname(os.path.abspath(__file__))
DEFAULT_VOICES = [
    {
        "label": "Elizabeth Klett (female)",
        "path": os.path.join(APP_DIR, "assets", "voices", "elizabeth_klett_female.mp3"),
    },
    {
        "label": "Mark F. Smith (male)",
        "path": os.path.join(APP_DIR, "assets", "voices", "mark_smith_male.mp3"),
    },
    {
        "label": "David Clarke (male)",
        "path": os.path.join(APP_DIR, "assets", "voices", "david_clarke_male.mp3"),
    },
]
DEFAULT_VOICE_CHOICES = [voice["label"] for voice in DEFAULT_VOICES]
DEFAULT_VOICE_PATH_BY_LABEL = {
    voice["label"]: voice["path"] for voice in DEFAULT_VOICES
}

PARALINGUISTIC_TAGS = [
    "[clear throat]",
    "[sigh]",
    "[shush]",
    "[cough]",
    "[groan]",
    "[sniff]",
    "[gasp]",
    "[chuckle]",
    "[laugh]",
]
PARALINGUISTIC_TAG_PATTERN = re.compile(r"\[[^\[\]]+\]")
PARALINGUISTIC_TAG_SET = set(PARALINGUISTIC_TAGS)
PRONUNCIATION_REPLACEMENTS = [
    (re.compile(r"\bzapply\b", flags=re.IGNORECASE), "zap-lee"),
]


def _env_flag(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def _resolve_device() -> str:
    override = os.getenv("CHATTERBOX_DEVICE")
    if override:
        return override
    return "cpu"


DEVICE = _resolve_device()
MODEL = ChatterboxTurboTTS.from_pretrained(device=DEVICE)


def _validate_script(script: str) -> str:
    script = (script or "").strip()
    if not script:
        raise gr.Error("Paste a script before generating a voiceover.")
    if len(script) > MAX_CHARS:
        raise gr.Error(f"Keep scripts to {MAX_CHARS} characters or fewer.")
    tags = PARALINGUISTIC_TAG_PATTERN.findall(script)
    unsupported_tags = sorted(set(tags) - PARALINGUISTIC_TAG_SET)
    if unsupported_tags:
        raise gr.Error(
            "Unsupported tag: "
            f"{', '.join(unsupported_tags)}. Use only {', '.join(PARALINGUISTIC_TAGS)}."
        )
    return script


def _apply_pronunciation_replacements(script: str) -> str:
    for pattern, replacement in PRONUNCIATION_REPLACEMENTS:
        script = pattern.sub(replacement, script)
    return script


def _save_wav(wav: torch.Tensor, sample_rate: int) -> str:
    if not isinstance(wav, torch.Tensor):
        wav = torch.as_tensor(wav)
    wav = wav.detach().cpu().float()
    if wav.ndim == 1:
        wav = wav.unsqueeze(0)

    output = tempfile.NamedTemporaryFile(
        delete=False,
        suffix=".wav",
        prefix="chatterbox_voiceover_",
    )
    output.close()
    ta.save(output.name, wav, sample_rate)
    return output.name


def generate_voiceover(
    script: str,
    reference_audio: Optional[str],
) -> Tuple[str, str]:
    script = _validate_script(script)
    script = _apply_pronunciation_replacements(script)
    generation_args = {}
    if reference_audio:
        generation_args["audio_prompt_path"] = reference_audio

    try:
        with torch.inference_mode():
            wav = MODEL.generate(script, **generation_args)
    except TypeError as exc:
        if not reference_audio:
            raise gr.Error(
                "Chatterbox Turbo requires a reference voice clip. "
                "Upload a short WAV, MP3, FLAC, M4A, or OGG file and try again."
            ) from exc
        raise gr.Error(f"Generation failed: {exc}") from exc
    except Exception as exc:
        raise gr.Error(f"Generation failed: {exc}") from exc

    output_path = _save_wav(wav, MODEL.sr)
    return output_path, output_path


def character_count(script: str) -> str:
    count = len(script or "")
    status = "OK" if count <= MAX_CHARS else "Too long"
    return f"{count}/{MAX_CHARS} characters - {status}"


def sync_script_and_count(script: str) -> Tuple[str, str]:
    return script, character_count(script)


def select_default_voice(default_voice: str) -> str:
    fallback_voice = DEFAULT_VOICE_CHOICES[0]
    return DEFAULT_VOICE_PATH_BY_LABEL.get(
        default_voice,
        DEFAULT_VOICE_PATH_BY_LABEL[fallback_voice],
    )


CSS = """
:root, body, .gradio-container {
  color-scheme: dark;
  background: #101114;
}
.gradio-container {
  color: #f4f4f5;
}
.tag-row button {
  min-width: 6rem;
}
.voiceover-output audio {
  width: 100%;
}
textarea, input, select {
  background: #14161a !important;
  color: #f4f4f5 !important;
}
"""

DARK_THEME = gr.themes.Soft().set(
    body_background_fill="#101114",
    body_background_fill_dark="#101114",
    body_text_color="#f4f4f5",
    body_text_color_dark="#f4f4f5",
    background_fill_primary="#101114",
    background_fill_primary_dark="#101114",
    background_fill_secondary="#17181c",
    background_fill_secondary_dark="#17181c",
    block_background_fill="#17181c",
    block_background_fill_dark="#17181c",
    block_border_color="#2a2d34",
    block_border_color_dark="#2a2d34",
    input_background_fill="#14161a",
    input_background_fill_dark="#14161a",
    input_border_color="#30333b",
    input_border_color_dark="#30333b",
    button_primary_background_fill="#7c3aed",
    button_primary_background_fill_dark="#7c3aed",
    button_primary_background_fill_hover="#8b5cf6",
    button_primary_background_fill_hover_dark="#8b5cf6",
    button_primary_text_color="#ffffff",
    button_primary_text_color_dark="#ffffff",
)

INSERT_TEXT_JS = """
(script) => {
  const insertText = INSERT_TEXT_VALUE;
  const active = document.activeElement;
  const textarea = active && active.tagName === "TEXTAREA"
    ? active
    : document.querySelector('textarea[data-testid="textbox"], textarea');

  if (!textarea) {
    const current = script || "";
    return [current + insertText];
  }

  const current = textarea.value || script || "";
  const start = textarea.selectionStart ?? current.length;
  const end = textarea.selectionEnd ?? start;
  const nextValue = current.slice(0, start) + insertText + current.slice(end);
  const nextCursor = start + insertText.length;

  textarea.value = nextValue;
  textarea.dispatchEvent(new Event("input", { bubbles: true }));
  requestAnimationFrame(() => {
    textarea.focus();
    textarea.setSelectionRange(nextCursor, nextCursor);
  });

  return [nextValue];
}
"""


def insert_text_js(insert_text: str) -> str:
    return INSERT_TEXT_JS.replace("INSERT_TEXT_VALUE", repr(insert_text))


with gr.Blocks(title="Chatterbox Voiceovers") as demo:
    gr.Markdown(
        """
        # Chatterbox Voiceovers
        Create short-form voiceovers from reel scripts with Chatterbox Turbo.
        Add expressive tags inline, choose a bundled voice or upload a voice
        reference, then generate a WAV.
        """
    )

    with gr.Row():
        with gr.Column(scale=7):
            script_input = gr.Textbox(
                label="Script",
                placeholder=(
                    "Paste your reel script here. Example: "
                    "This one mistake is costing creators hours every week [chuckle]."
                ),
                lines=12,
                max_lines=18,
            )
            count_output = gr.Markdown(character_count(""))

            with gr.Group(elem_classes=["tag-row"]):
                gr.Markdown("Paralinguistic tags")
                with gr.Row():
                    tag_buttons = [gr.Button(tag, size="sm") for tag in PARALINGUISTIC_TAGS]

        with gr.Column(scale=5):
            default_voice = gr.Dropdown(
                label="Default voice",
                choices=DEFAULT_VOICE_CHOICES,
                value=DEFAULT_VOICE_CHOICES[0],
            )
            reference_audio = gr.Audio(
                label="Voice reference",
                sources=["upload", "microphone"],
                type="filepath",
                value=DEFAULT_VOICE_PATH_BY_LABEL[DEFAULT_VOICE_CHOICES[0]],
            )
            generate_button = gr.Button("Generate voiceover", variant="primary")

            output_audio = gr.Audio(
                label="Generated voiceover",
                type="filepath",
                elem_classes=["voiceover-output"],
            )
            download_file = gr.File(label="Download WAV")

    gr.Examples(
        examples=[
            [
                "Stop scrolling [chuckle]. If your reels sound flat, try writing like you speak, then cut every sentence in half."
            ],
            [
                "I thought this would take all afternoon [sigh], but the whole edit was done before my coffee got cold."
            ],
            [
                "Here is the simple three-part hook that keeps people watching [chuckle]. Problem, surprise, payoff."
            ],
        ],
        inputs=[script_input],
    )

    script_input.change(character_count, inputs=script_input, outputs=count_output)
    default_voice.change(select_default_voice, inputs=default_voice, outputs=reference_audio)
    for tag_button, tag in zip(tag_buttons, PARALINGUISTIC_TAGS):
        tag_button.click(
            fn=sync_script_and_count,
            inputs=script_input,
            outputs=[script_input, count_output],
            js=insert_text_js(f"{tag} "),
        )
    generate_button.click(
        fn=generate_voiceover,
        inputs=[script_input, reference_audio],
        outputs=[output_audio, download_file],
        show_progress="full",
    )


if __name__ == "__main__":
    demo.queue(max_size=20, default_concurrency_limit=1).launch(
        theme=DARK_THEME,
        css=CSS,
        share=_env_flag("GRADIO_SHARE", default=False),
        ssr_mode=False,
    )