import time
import solara
import soundfile as sf
from kokoro_onnx import Kokoro
from IPython.display import display, Audio

text = solara.reactive("")

@solara.component
def Example(example: str, example_shortened=None):
    def on_click():
        text.value = example
    with solara.Row():
        solara.Button(example, on_click=on_click)


models = ["kokoro-v0_19.onnx", "kokoro-quant.onnx"]
model = solara.reactive("kokoro-v0_19.onnx")
voices = ["af", "af_bella", "af_nicole", "af_sarah", "af_sky", "am_adam", "am_michael", "bf_emma", "bf_isabella", "bm_george", "bm_lewis"]
voice = solara.reactive("af_sarah")

generate = solara.reactive(False)
@solara.component
def Page():
    title = "Kokoro Text-To-Speech"
    with solara.Head():
        solara.Title(f"{title}")
    with solara.Column(style={"width": "100%", "padding": "50px", "align": "center"}):
        solara.Markdown(f"#{title}")
        solara.Markdown("Try some examples:")
        with solara.Row():
            Example("Hello! How can I assist you today?")
            Example("The capital of France is Paris.")
        solara.Markdown("Or write your own:")
        solara.InputTextArea("Enter some text", value=text, rows=3, continuous_update=False)
        with solara.Row():
            solara.Select(label="Select voice:", value=voice, values=voices, style={"width": "10%"})
            solara.Select(label="Select model:", value=model, values=models, style={"width": "10%"})
            solara.Button("Generate", on_click=generate.set(True))
            solara.Button("Clear", on_click=lambda: text.set(""))
        kokoro = Kokoro(model.value, "voices.bin")
        if text.value != "" and generate.value:
            start_time = time.time()
            samples, sample_rate = kokoro.create(text.value, voice=voice.value, speed=1.0)
            end_time = time.time()
            display(Audio(data=samples, rate=sample_rate, autoplay=True))
            solara.Text(f"Generation time: {end_time - start_time:.1f} seconds", style="color: blue")