import re
import sys

import gradio as gr
import yaml
from huggingface_hub import snapshot_download

try:
    import spaces

    USING_SPACES = True
except ImportError:
    USING_SPACES = False

MODEL_ID = "ILRDF/chatterbox-formosan-all-finetune"

repo_dir = snapshot_download(MODEL_ID)
sys.path.insert(0, repo_dir)
from formosan_tts import FormosanTTS  # noqa: E402

tts = FormosanTTS(repo_dir, device="cuda")


def gpu_decorator(func):
    if USING_SPACES:
        return spaces.GPU(func)
    else:
        return func


with open("configs/refs.yaml", encoding="utf-8") as f:
    refs_config = yaml.safe_load(f)

# "阿美_海岸" -> "ami-x-pswl"; ethnicity "阿美" -> ["阿美_海岸", ...]
LANGUAGES = {}
for ref, v in refs_config.items():
    LANGUAGES.setdefault(re.sub(r"_[男女]聲[12]?$", "", ref), v["lang_code"])
ETHNICITIES = sorted({lang.split("_")[0] for lang in LANGUAGES})


def languages_of(ethnicity):
    return [lang for lang in LANGUAGES if lang.split("_")[0] == ethnicity]


def refs_of(prefix):
    return [r for r in refs_config if r.startswith(prefix)]


@gpu_decorator
def synthesize(gen_text, lang_code, ref_audio):
    gen_text = gen_text.strip()
    if len(gen_text) == 0:
        raise gr.Error("請勿輸入空字串。")
    if not ref_audio:
        raise gr.Error("請先錄製或上傳參考音檔。")
    try:
        wav, sr = tts.generate(gen_text, lang_code, ref_audio)
    except ValueError as e:
        raise gr.Error(f"文字中有無法轉換的字元：{e}")
    return sr, wav


def get_title():
    with open("DEMO.md", encoding="utf-8") as tong:
        return tong.readline().strip("# ")


CSS = """@import url(https://tauhu.tw/tauhu-oo.css);
.textonly textarea {border-width: 0px !important; }
"""

THEME = gr.themes.Default(
    font=(
        "tauhu-oo",
        gr.themes.GoogleFont("Source Sans Pro"),
        "ui-sans-serif",
        "system-ui",
        "sans-serif",
    )
)

JS = """
function addButtonsEvent() {
    const buttons = document.querySelectorAll("#head-html-block button");
    buttons.forEach(button => {
        button.addEventListener("click", () => {
            navigator.clipboard.writeText(button.innerText);
        });
    });
}
"""

demo = gr.Blocks(title=get_title())

with demo:
    with open("DEMO.md", encoding="utf-8") as tong:
        gr.Markdown(tong.read())

    gr.HTML(
        "特殊符號請複製使用（滑鼠點擊即可複製）：<button>é</button> <button>ṟ</button> <button>ɨ</button> <button>ʉ</button>",
        padding=False,
        elem_id="head-html-block",
    )

    with gr.Tab("預設配音員"):
        with gr.Row():
            with gr.Column():
                default_speaker_ethnicity = gr.Dropdown(
                    choices=ETHNICITIES,
                    label="步驟一：選擇族別",
                    value="阿美",
                    filterable=False,
                )

                default_speaker_refs = gr.Dropdown(
                    choices=refs_of(default_speaker_ethnicity.value),
                    label="步驟二：選擇配音員",
                    value=refs_of(default_speaker_ethnicity.value)[0],
                    filterable=False,
                )

                default_speaker_gen_text_input = gr.Textbox(
                    label="步驟三：輸入文字",
                    value="",
                )

                default_speaker_generate_btn = gr.Button(
                    "步驟四：開始合成", variant="primary"
                )

            with gr.Column():
                default_speaker_audio_output = gr.Audio(
                    label="合成結果", buttons=["download"]
                )

    with gr.Tab("自己當配音員"):
        with gr.Row():
            with gr.Column():
                custom_speaker_ethnicity = gr.Dropdown(
                    choices=ETHNICITIES,
                    label="步驟一：選擇族別與語別",
                    value="阿美",
                    filterable=False,
                )

                custom_speaker_language = gr.Dropdown(
                    choices=languages_of(custom_speaker_ethnicity.value),
                    value=languages_of(custom_speaker_ethnicity.value)[0],
                    filterable=False,
                    show_label=False,
                )

                custom_speaker_ref_text = gr.Textbox(
                    value=refs_config[refs_of(custom_speaker_language.value)[0]]["text"],
                    interactive=False,
                    label="步驟二：點選🎙️錄製一段清楚的聲音（建議 3 秒以上，可念下方句子），或上傳音檔",
                    elem_classes="textonly",
                )

                custom_speaker_audio_input = gr.Audio(
                    type="filepath",
                    sources=["microphone", "upload"],
                    waveform_options=gr.WaveformOptions(
                        sample_rate=24000,
                    ),
                    label="錄製或上傳",
                )

                custom_speaker_gen_text_input = gr.Textbox(
                    label="步驟三：輸入合成文字",
                    value="",
                )

                custom_speaker_generate_btn = gr.Button(
                    "步驟四：開始合成", variant="primary"
                )

            with gr.Column():
                custom_speaker_audio_output = gr.Audio(
                    label="合成結果", buttons=["download"]
                )

    default_speaker_ethnicity.change(
        lambda ethnicity: gr.Dropdown(
            choices=refs_of(ethnicity),
            value=refs_of(ethnicity)[0],
        ),
        inputs=[default_speaker_ethnicity],
        outputs=[default_speaker_refs],
    )

    def default_speaker_tts(ref, gen_text_input):
        return synthesize(gen_text_input, refs_config[ref]["lang_code"], refs_config[ref]["wav"])

    default_speaker_generate_btn.click(
        default_speaker_tts,
        inputs=[default_speaker_refs, default_speaker_gen_text_input],
        outputs=[default_speaker_audio_output],
    )

    custom_speaker_ethnicity.change(
        lambda ethnicity: gr.Dropdown(
            choices=languages_of(ethnicity),
            value=languages_of(ethnicity)[0],
            visible=len(languages_of(ethnicity)) > 1,
        ),
        inputs=[custom_speaker_ethnicity],
        outputs=[custom_speaker_language],
    )

    custom_speaker_language.change(
        lambda lang: gr.Textbox(value=refs_config[refs_of(lang)[0]]["text"]),
        inputs=[custom_speaker_language],
        outputs=[custom_speaker_ref_text],
    )

    def custom_speaker_tts(language, ref_audio_input, gen_text_input):
        return synthesize(gen_text_input, LANGUAGES[language], ref_audio_input)

    custom_speaker_generate_btn.click(
        custom_speaker_tts,
        inputs=[custom_speaker_language, custom_speaker_audio_input, custom_speaker_gen_text_input],
        outputs=[custom_speaker_audio_output],
    )

demo.launch(css=CSS, theme=THEME, js=JS)