import gradio as gr
import numpy as np
import base64
from pathlib import Path

SAMPLE_RATE = 24000
ROOT_DIR = Path(__file__).parent
LOGO_PATH = ROOT_DIR / "assets" / "kittenml_logo.svg"

MODELS = {
    "Nano (15M - Fastest)": "KittenML/kitten-tts-nano-0.8-fp32",
    "Micro (40M - Balanced)": "KittenML/kitten-tts-micro-0.8",
    "Mini (80M - Best Quality)": "KittenML/kitten-tts-mini-0.8",
}

VOICES = [
    "Bella",
    "Jasper",
    "Luna",
    "Bruno",
    "Rosie",
    "Hugo",
    "Kiki",
    "Leo",
]

EXAMPLES = [
    (
        "Space is a three-dimensional continuum containing positions and directions.",
        "Micro (40M - Balanced)",
        "Jasper",
        1.0,
    ),
    (
        "She picked up her coffee and walked toward the window.",
        "Mini (80M - Best Quality)",
        "Luna",
        1.0,
    ),
    (
        "The sun set slowly over the calm, quiet lake",
        "Nano (15M - Fastest)",
        "Bella",
        1.1,
    ),
]

_model_cache: dict[str, object] = {}


def logo_data_uri() -> str:
    if not LOGO_PATH.exists():
        return "https://kittenml.com/assets/kittenml_logo.svg"

    encoded_logo = base64.b64encode(LOGO_PATH.read_bytes()).decode("ascii")
    return f"data:image/svg+xml;base64,{encoded_logo}"


def get_model(model_name: str):
    if model_name not in _model_cache:
        from kittentts import KittenTTS

        print(f"Loading {model_name}...")
        _model_cache[model_name] = KittenTTS(MODELS[model_name])
        print(f"{model_name} loaded.")

    return _model_cache[model_name]


def synthesize(text: str, model_name: str, voice: str, speed: float):
    if not text or not text.strip():
        raise gr.Error("Please enter some text.")

    tts = get_model(model_name)
    # Note: speed parameter may not be supported in v0.8
    # If you get an error, remove speed=speed from the generate call
    try:
        audio = tts.generate(text.strip(), voice=voice, speed=speed)
    except TypeError:
        # Fallback if speed is not supported
        audio = tts.generate(text.strip(), voice=voice)

    # audio shape is (1, samples) or (samples,) — normalize to 1-D
    audio = np.squeeze(audio)

    return (SAMPLE_RATE, audio)


theme = gr.themes.Base(
    primary_hue="neutral",
    secondary_hue="neutral",
    neutral_hue="neutral",
    font=gr.themes.GoogleFont("Inter"),
).set(
    body_background_fill="#f7f7f8",
    body_background_fill_dark="#f7f7f8",
    block_background_fill="#ffffff",
    block_background_fill_dark="#ffffff",
    block_border_color="#e7e7e8",
    block_border_color_dark="#e7e7e8",
    block_radius="8px",
    block_shadow="0 1px 2px rgba(0, 0, 0, 0.04)",
    block_shadow_dark="0 1px 2px rgba(0, 0, 0, 0.04)",
    button_primary_background_fill="#101010",
    button_primary_background_fill_hover="#18181b",
    button_primary_text_color="white",
    button_primary_border_color="#101010",
    input_background_fill="white",
    input_background_fill_dark="white",
    input_border_color="transparent",
    input_border_color_focus="#101010",
    slider_color="#101010",
    table_border_color="transparent",
    table_even_background_fill="white",
    table_odd_background_fill="#fafafa",
    table_row_focus="white",
)

css = """
:root, html, body { color-scheme: light !important; }
body {
    background: #f7f7f8 !important;
}
.gradio-container,
.main,
.app,
.contain,
.wrap,
.gap,
.form {
    background: transparent !important;
}
.gradio-container {
    max-width: 1000px !important;
    margin: 0 auto !important;
    padding: 22px 18px 30px !important;
    --section-inset: 12px;
}
footer { display: none !important; }

.hero {
    position: relative;
    overflow: hidden;
    width: 100%;
    max-width: 100%;
    box-sizing: border-box;
    border: 1px solid rgba(150, 150, 156, 0.42);
    border-radius: 8px;
    background:
        linear-gradient(116deg, rgba(255, 255, 255, 0.78) 0%, rgba(255, 255, 255, 0) 17%, rgba(255, 255, 255, 0.55) 34%, rgba(255, 255, 255, 0) 51%),
        linear-gradient(145deg, #f8f8f9 0%, #dedee3 31%, #fdfdfd 54%, #bfc0c7 100%);
    display: flex;
    flex-direction: column;
    justify-content: space-between;
    min-height: 266px;
    padding: 38px 40px 34px;
    box-shadow:
        inset 0 1px 0 rgba(255, 255, 255, 0.82),
        inset 0 -1px 0 rgba(24, 24, 27, 0.08),
        0 22px 54px rgba(24, 24, 27, 0.12),
        0 1px 2px rgba(24, 24, 27, 0.08);
}
.hero::before {
    background:
        repeating-linear-gradient(90deg, rgba(255, 255, 255, 0.16) 0, rgba(255, 255, 255, 0.16) 1px, rgba(35, 35, 39, 0.025) 2px, rgba(35, 35, 39, 0.025) 4px),
        linear-gradient(90deg, rgba(17, 17, 17, 0.11), rgba(17, 17, 17, 0) 42%);
    content: "";
    inset: 0;
    opacity: 0.55;
    pointer-events: none;
    position: absolute;
}
.hero::after {
    background: linear-gradient(180deg, rgba(255, 255, 255, 0.38), rgba(255, 255, 255, 0) 48%, rgba(24, 24, 27, 0.035));
    content: "";
    inset: 0;
    pointer-events: none;
    position: absolute;
}
.brand-row {
    display: flex;
    align-items: center;
    justify-content: space-between;
    gap: 16px;
    margin-bottom: 38px;
    position: relative;
    z-index: 2;
}
.brand-lockup {
    align-items: center;
    background: #ffffff;
    border-radius: 999px;
    border: 1px solid rgba(255, 255, 255, 0.92);
    box-shadow:
        inset 0 1px 0 rgba(255, 255, 255, 0.88),
        0 12px 28px rgba(24, 24, 27, 0.10),
        0 1px 2px rgba(24, 24, 27, 0.08);
    display: inline-flex;
    gap: 11px;
    padding: 9px 15px 9px 10px;
}
.brand-mark {
    height: 34px;
    width: auto;
    display: block;
}
.brand-lockup span {
    color: #111111;
    font-size: 13px;
    font-weight: 750;
}
.status-pill {
    border: 0;
    border-radius: 999px;
    color: #111111;
    font-size: 12px;
    font-weight: 700;
    padding: 9px 14px;
    background: #ffffff;
    border: 1px solid rgba(255, 255, 255, 0.92);
    box-shadow:
        inset 0 1px 0 rgba(255, 255, 255, 0.88),
        0 12px 28px rgba(24, 24, 27, 0.10),
        0 1px 2px rgba(24, 24, 27, 0.08);
}
.hero h1 {
    background: linear-gradient(180deg, #050506 0%, #2f3035 48%, #070708 100%);
    -webkit-background-clip: text;
    background-clip: text;
    color: transparent;
    font-size: clamp(40px, 5.4vw, 66px);
    font-weight: 820;
    letter-spacing: 0;
    line-height: 0.96;
    margin: 0;
    max-width: 780px;
    position: relative;
    text-shadow: 0 1px 0 rgba(255, 255, 255, 0.42);
    z-index: 2;
}
.spec-grid {
    display: flex;
    flex-wrap: wrap;
    gap: 7px;
    margin-top: 22px;
    max-width: 520px;
    position: relative;
    z-index: 2;
}
.spec {
    align-items: center;
    background: rgba(255, 255, 255, 0.68);
    border: 1px solid rgba(255, 255, 255, 0.76);
    border-radius: 999px;
    box-shadow: inset 0 1px 0 rgba(255, 255, 255, 0.64), 0 8px 18px rgba(24, 24, 27, 0.08);
    display: inline-flex;
    padding: 7px 12px;
}
.spec strong {
    color: #111111;
    font-size: 12px;
    font-weight: 700;
    line-height: 1;
}
.studio-layout {
    align-items: stretch !important;
    gap: 14px !important;
    margin: 14px var(--section-inset) 0 !important;
    width: calc(100% - (var(--section-inset) * 2)) !important;
    max-width: calc(100% - (var(--section-inset) * 2)) !important;
    box-sizing: border-box;
}
.studio-layout.unequal-height {
    align-items: stretch !important;
}
.studio-card {
    align-self: stretch !important;
    border: 1px solid #e7e7e8 !important;
    border-radius: 8px !important;
    background: #ffffff !important;
    box-sizing: border-box !important;
    box-shadow: 0 1px 2px rgba(0, 0, 0, 0.04) !important;
    padding: 16px !important;
}
.gradio-container {
    --body-text-color: #111111 !important;
    --block-label-text-color: #2c2924 !important;
    --block-title-text-color: #2c2924 !important;
    --input-text-color: #111111 !important;
}
.gradio-container .block,
.gradio-container .panel {
    background: transparent !important;
    border-color: transparent !important;
    border-width: 0 !important;
    box-shadow: none !important;
    color: #111111 !important;
}
.gradio-container .block,
.gradio-container .block > *,
.gradio-container .form,
.gradio-container .wrap,
.gradio-container .panel,
.gradio-container .gap {
    box-shadow: none !important;
}
.gradio-container .block *,
.gradio-container .form * {
    color: #111111 !important;
}
.gradio-container label,
.gradio-container label span,
.gradio-container .block-label,
.gradio-container .block-title,
.gradio-container [data-testid="block-label"],
.gradio-container [data-testid="block-label"] * {
    color: #2c2924 !important;
    font-weight: 650 !important;
}
.gradio-container [data-testid="block-label"] {
    background: transparent !important;
    border-color: transparent !important;
}
.gradio-container textarea,
.gradio-container .tab-like-container {
    background: #f7f7f7 !important;
    border-color: transparent !important;
    border-width: 0 !important;
    border-radius: 8px !important;
    box-shadow: none !important;
}
.gradio-container .wrap-inner {
    background: transparent !important;
    border-radius: 8px !important;
    box-shadow: none !important;
}
.gradio-container .secondary-wrap,
.gradio-container .wrap-inner > *,
.gradio-container input,
.gradio-container input[role="listbox"] {
    background: transparent !important;
    border-color: transparent !important;
    border-width: 0 !important;
    box-shadow: none !important;
}
.gradio-container textarea {
    min-height: 128px !important;
}
.gradio-container input[role="listbox"] {
    background: #f7f7f7 !important;
    border-radius: 8px !important;
    min-height: 42px !important;
    padding: 0 14px !important;
}
.gradio-container textarea, .gradio-container input {
    color: #111111 !important;
}
::placeholder { color: #7f7f7f !important; }
.generate-button button {
    min-height: 46px !important;
    border-radius: 8px !important;
    font-weight: 750 !important;
    box-shadow: none !important;
}
.generate-button button,
.generate-button button *,
button.primary,
button.primary * {
    color: #ffffff !important;
}
.gradio-container button:hover {
    background: #eeeeee !important;
    color: #111111 !important;
}
.gradio-container button:hover * {
    color: #111111 !important;
}
.generate-button button:hover,
.generate-button button:hover *,
button.primary:hover,
button.primary:hover * {
    background: #18181b !important;
    color: #ffffff !important;
}
.gradio-container ul.options,
.gradio-container .options {
    background: #ffffff !important;
    border: 1px solid #e5e5e5 !important;
    border-radius: 8px !important;
    box-shadow: 0 14px 34px rgba(0, 0, 0, 0.12) !important;
    overflow: hidden !important;
}
.gradio-container .wrap-inner.show_options,
.gradio-container .wrap-inner.show_options * {
    background: #ffffff !important;
    color: #111111 !important;
}
.gradio-container ul.options li,
.gradio-container .options li,
.gradio-container [role="option"] {
    background: #ffffff !important;
    color: #111111 !important;
}
.gradio-container ul.options li:hover,
.gradio-container ul.options li.selected,
.gradio-container .options li:hover,
.gradio-container .options li.selected,
.gradio-container [role="option"]:hover,
.gradio-container [role="option"][aria-selected="true"] {
    background: #f0f0f0 !important;
    color: #111111 !important;
}
.output-panel {
    justify-content: space-between !important;
    min-height: 100%;
}
.output-panel audio,
.output-panel [data-testid="audio"] {
    min-height: 184px !important;
}
.output-panel .progress-text.meta-text {
    align-items: center !important;
    background: #f4f4f5 !important;
    border: 1px solid #e4e4e7 !important;
    border-radius: 999px !important;
    color: transparent !important;
    display: inline-flex !important;
    font-family: Inter, Arial, sans-serif !important;
    font-size: 0 !important;
    font-weight: 650 !important;
    justify-content: center !important;
    line-height: 1 !important;
    margin: 10px auto 0 !important;
    min-width: 0 !important;
    padding: 7px 10px !important;
    position: static !important;
    text-align: center !important;
    width: fit-content !important;
}
.output-panel .progress-text.meta-text::before {
    color: #3f3f46;
    content: "Generating audio";
    font-size: 12px;
    letter-spacing: 0;
}
.output-note {
    border: 0;
    border-radius: 8px;
    margin-top: 16px;
    padding: 14px 15px;
    background: #f7f7f7;
    color: #3f3f3f;
    font-size: 13px;
    line-height: 1.5;
}
.output-note strong {
    color: #151515;
    display: block;
    margin-bottom: 4px;
}
.examples-panel {
    margin: 14px var(--section-inset) 0 !important;
    padding: 14px !important;
    width: calc(100% - (var(--section-inset) * 2)) !important;
    max-width: calc(100% - (var(--section-inset) * 2)) !important;
    box-sizing: border-box;
    background: #ffffff !important;
    border: 1px solid #e5e5e5 !important;
    border-radius: 8px !important;
    box-shadow: 0 1px 2px rgba(0, 0, 0, 0.04) !important;
    overflow: hidden !important;
}
.examples-header {
    display: flex;
    align-items: center;
    justify-content: space-between;
    gap: 14px;
    margin-bottom: 10px;
}
.examples-header strong {
    color: #111111;
    font-size: 15px;
}
.examples-header span {
    color: #666666;
    font-size: 12px;
}
.example-row {
    align-items: center !important;
    background: #f8f8f8 !important;
    border: 1px solid #eeeeee !important;
    border-radius: 8px !important;
    margin-top: 8px !important;
    padding: 10px 10px 10px 12px !important;
    transition: background 120ms ease, border-color 120ms ease;
}
.example-row > .block {
    min-width: 0 !important;
}
.example-row:hover {
    background: #f4f4f5 !important;
    border-color: #e2e2e3 !important;
}
.example-copy {
    display: grid;
    grid-template-columns: minmax(0, 1fr) 190px 80px 52px;
    gap: 14px;
    align-items: center;
    width: 100%;
}
.example-copy strong,
.example-copy span {
    color: #111111;
    display: block;
    overflow: hidden;
    text-overflow: ellipsis;
    white-space: nowrap;
}
.example-copy strong {
    font-size: 13px;
    font-weight: 650;
}
.example-copy span {
    color: #5c5c5c;
    font-size: 12px;
}
.load-example {
    flex: 0 0 72px !important;
    max-width: 72px !important;
    min-width: 72px !important;
    width: 72px !important;
}
.load-example {
    align-items: center !important;
    background: #111111 !important;
    border: 1px solid #111111 !important;
    border-radius: 8px !important;
    color: #ffffff !important;
    display: flex !important;
    font-size: 13px !important;
    font-weight: 750 !important;
    height: 40px !important;
    justify-content: center !important;
    line-height: 1 !important;
    min-width: 72px !important;
    padding: 0 !important;
    width: 72px !important;
}
.gradio-container button.load-example,
button.load-example {
    background: #111111 !important;
    border-color: #111111 !important;
    color: #ffffff !important;
}
.gradio-container button.load-example *,
button.load-example * {
    background: transparent !important;
    color: #ffffff !important;
}
.gradio-container button.load-example:hover,
button.load-example:hover {
    background: #18181b !important;
    border-color: #18181b !important;
    color: #ffffff !important;
}
.gradio-container button.load-example:hover *,
button.load-example:hover * {
    background: transparent !important;
    color: #ffffff !important;
}

textarea,
input,
select,
table {
    border-width: 0 !important;
    box-shadow: none !important;
}

@media (max-width: 760px) {
    .gradio-container { padding: 16px 12px 24px !important; }
    .hero { min-height: 230px; padding: 22px; }
    .brand-row { align-items: flex-start; flex-direction: column; margin-bottom: 24px; }
    .hero h1 { font-size: clamp(38px, 13vw, 54px); }
    .spec-grid { gap: 7px; }
    .spec { max-width: 100%; }
    .example-copy { grid-template-columns: 1fr; gap: 4px; }
}
"""

hero_html = f"""
<section class="hero">
    <div class="brand-row">
        <div class="brand-lockup">
            <img class="brand-mark" src="{logo_data_uri()}" alt="KittenML">
            <span>KittenML</span>
        </div>
        <div class="status-pill">On-device friendly TTS</div>
    </div>
    <h1>KittenTTS Demo</h1>
    <div class="spec-grid">
        <div class="spec"><strong>24 kHz audio</strong></div>
        <div class="spec"><strong>8 voices</strong></div>
        <div class="spec"><strong>15M-80M models</strong></div>
    </div>
</section>
"""

with gr.Blocks(title="KittenTTS Voice Studio") as demo:
    gr.HTML(hero_html)

    with gr.Row(elem_classes=["studio-layout"]):
        with gr.Column(scale=2, elem_classes=["studio-card"]):
            text_input = gr.Textbox(
                label="Script",
                placeholder="Write a short passage for KittenTTS to read aloud.",
                lines=4,
                max_lines=6,
            )
            with gr.Row():
                model_select = gr.Dropdown(
                    choices=list(MODELS.keys()),
                    value="Micro (40M - Balanced)",
                    label="Model",
                )
                voice_select = gr.Dropdown(
                    choices=VOICES,
                    value="Jasper",
                    label="Voice",
                )
            speed_slider = gr.Slider(
                minimum=0.5,
                maximum=2.0,
                value=1.0,
                step=0.05,
                label="Pace",
            )
            generate_btn = gr.Button(
                "Generate speech",
                variant="primary",
                elem_classes=["generate-button"],
            )

        with gr.Column(scale=1, elem_classes=["studio-card", "output-panel"]):
            audio_output = gr.Audio(label="Generated audio", type="numpy")
            gr.HTML(
                """
                <div class="output-note">
                    <strong>Model loading</strong>
                    The selected model loads on first use, then stays cached for faster generations.
                </div>
                """
            )

    generate_btn.click(
        fn=synthesize,
        inputs=[text_input, model_select, voice_select, speed_slider],
        outputs=audio_output,
    )

    with gr.Column(elem_classes=["examples-panel"]):
        gr.HTML(
            """
            <div class="examples-header">
                <strong>Examples</strong>
                <span>Load a sample into the studio</span>
            </div>
            """
        )
        for text, model, voice, pace in EXAMPLES:
            with gr.Row(elem_classes=["example-row"]):
                gr.HTML(
                    f"""
                    <div class="example-copy">
                        <strong>{text}</strong>
                        <span>{model}</span>
                        <span>{voice}</span>
                        <span>{pace:g}x</span>
                    </div>
                    """
                )
                load_btn = gr.Button("Load", elem_classes=["load-example"])
                load_btn.click(
                    fn=lambda text=text, model=model, voice=voice, pace=pace: (
                        text,
                        model,
                        voice,
                        pace,
                    ),
                    inputs=[],
                    outputs=[text_input, model_select, voice_select, speed_slider],
                )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", theme=theme, css=css)