import base64
import gradio as gr
import librosa
import logging
import os
import soundfile as sf
import subprocess
import tempfile
import urllib.request

from datetime import datetime
from time import time

from examples import examples
from model import UETASRModel


def get_duration(filename: str) -> float:
    return librosa.get_duration(path=filename)


def convert_to_wav(in_filename: str) -> str:
    out_filename = os.path.splitext(in_filename)[0] + ".wav"
    logging.info(f"Converting {in_filename} to {out_filename}")
    y, sr = librosa.load(in_filename, sr=16000)
    sf.write(out_filename, y, sr)
    return out_filename


def build_html_output(s: str, style: str = "result_item_success"):
    return f"""
    <div class='result'>
        <div class='result_item {style}'>
          {s}
        </div>
    </div>
    """


def process_url(
    url: str,
    decoding_method: str,
    beam_size: int,
    max_symbols_per_step: int,
):
    logging.info(f"Processing URL: {url}")
    with tempfile.NamedTemporaryFile() as f:
        try:
            urllib.request.urlretrieve(url, f.name)
            return process(in_filename=f.name,
                           decoding_method=decoding_method,
                           beam_size=beam_size,
                           max_symbols_per_step=max_symbols_per_step)
        except Exception as e:
            logging.info(str(e))
            return "", build_html_output(str(e), "result_item_error")


def process_uploaded_file(
    in_filename: str,
    decoding_method: str,
    beam_size: int,
    max_symbols_per_step: int,
):
    if in_filename is None or in_filename == "":
        return "", build_html_output(
            "Please first upload a file and then click "
            'the button "submit for recognition"',
            "result_item_error",
        )

    logging.info(f"Processing uploaded file: {in_filename}")
    try:
        return process(in_filename=in_filename,
                       decoding_method=decoding_method,
                       beam_size=beam_size,
                       max_symbols_per_step=max_symbols_per_step)
    except Exception as e:
        logging.info(str(e))
        return "", build_html_output(str(e), "result_item_error")


def process_microphone(
    in_filename: str,
    decoding_method: str,
    beam_size: int,
    max_symbols_per_step: int,
):
    if in_filename is None or in_filename == "":
        return "", build_html_output(
            "Please first upload a file and then click "
            'the button "submit for recognition"',
            "result_item_error",
        )

    logging.info(f"Processing microphone: {in_filename}")
    try:
        return process(in_filename=in_filename,
                       decoding_method=decoding_method,
                       beam_size=beam_size,
                       max_symbols_per_step=max_symbols_per_step)
    except Exception as e:
        logging.info(str(e))
        return "", build_html_output(str(e), "result_item_error")


def process(
    in_filename: str,
    decoding_method: str,
    beam_size: int,
    max_symbols_per_step: int,
):
    logging.info(f"in_filename: {in_filename}")

    filename = convert_to_wav(in_filename)

    now = datetime.now()
    date_time = now.strftime("%d/%m/%Y, %H:%M:%S.%f")
    logging.info(f"Started at {date_time}")

    repo_id = "thanhtvt/uetasr-conformer_30.3m"

    start = time()

    recognizer = UETASRModel(repo_id,
                             decoding_method,
                             beam_size,
                             max_symbols_per_step)
    text = recognizer.predict(filename)

    date_time = now.strftime("%d/%m/%Y, %H:%M:%S.%f")
    end = time()

    duration = get_duration(filename)
    rtf = (end - start) / duration

    logging.info(f"Finished at {date_time} s. Elapsed: {end - start: .3f} s")

    info = f"""
    Wave duration  : {duration: .3f} s <br/>
    Processing time: {end - start: .3f} s <br/>
    RTF: {end - start: .3f}/{duration: .3f} = {rtf:.3f} <br/>
    """
    if rtf > 1:
        info += (
            "<br/>We are loading required resources for the first run. "
            "Please run again to measure the real RTF.<br/>"
        )

    logging.info(info)

    return text, build_html_output(info)


title = "Vietnamese Automatic Speech Recognition with UETASR"
description = """
This space shows how to use UETASR for Vietnamese Automatic Speech Recognition.

It is running on CPU provided by Hugging Face 🤗

See more information by visiting the [Github repository](https://github.com/thanhtvt/uetasr/)
"""

# css style is copied from
# https://huggingface.co/spaces/alphacep/asr/blob/main/app.py#L113
css = """
.result {display:flex;flex-direction:column}
.result_item {padding:15px;margin-bottom:8px;border-radius:15px;width:100%}
.result_item_success {background-color:mediumaquamarine;color:white;align-self:start}
.result_item_error {background-color:#ff7070;color:white;align-self:start}
"""

demo = gr.Blocks(css=css)


with demo:
    gr.Markdown(title)

    decode_method_radio = gr.Radio(
        label="Decoding method",
        choices=["greedy_search", "beam_search"],
        value="greedy_search",
        interactive=True,
    )

    beam_size_slider = gr.Slider(
        label="Beam size",
        minimum=1,
        maximum=20,
        step=1,
        value=1,
        interactive=False,
    )

    def interact_beam_slider(decoding_method):
        if decoding_method == "greedy_search":
            return gr.update(value=1, interactive=False)
        else:
            return gr.update(interactive=True)

    decode_method_radio.change(interact_beam_slider,
                               decode_method_radio,
                               beam_size_slider)

    max_symbols_per_step_slider = gr.Slider(
        label="Maximum symbols per step",
        minimum=1,
        maximum=20,
        step=1,
        value=5,
        interactive=True,
        visible=True,
    )

    with gr.Tabs():
        with gr.TabItem("Upload from disk"):
            uploaded_file = gr.Audio(
                source="upload",  # Choose between "microphone", "upload"
                type="filepath",
                label="Upload from disk",
            )
            upload_button = gr.Button("Submit for recognition")
            uploaded_output = gr.Textbox(label="Recognized speech from uploaded file")
            uploaded_html_info = gr.HTML(label="Info")

            gr.Examples(
                examples=examples,
                inputs=uploaded_file,
                outputs=[uploaded_output, uploaded_html_info],
                fn=process_uploaded_file,
            )

        with gr.TabItem("Record from microphone"):
            microphone = gr.Audio(
                source="microphone",
                type="filepath",
                label="Record from microphone",
            )

            record_button = gr.Button("Submit for recognition")
            recorded_output = gr.Textbox(label="Recognized speech from recordings")
            recorded_html_info = gr.HTML(label="Info")

            gr.Examples(
                examples=examples,
                inputs=microphone,
                outputs=[uploaded_output, uploaded_html_info],
                fn=process_microphone,
            )

        with gr.TabItem("From URL"):
            url_textbox = gr.Textbox(
                max_lines=1,
                placeholder="URL to an audio file",
                label="URL",
                interactive=True,
            )

            url_button = gr.Button("Submit for recognition")
            url_output = gr.Textbox(label="Recognized speech from URL")
            url_html_info = gr.HTML(label="Info")

        upload_button.click(
            process_uploaded_file,
            inputs=[
                uploaded_file,
                decode_method_radio,
                beam_size_slider,
                max_symbols_per_step_slider,
            ],
            outputs=[uploaded_output, uploaded_html_info],
        )

        record_button.click(
            process_microphone,
            inputs=[
                microphone,
                decode_method_radio,
                beam_size_slider,
                max_symbols_per_step_slider,
            ],
            outputs=[recorded_output, recorded_html_info],
        )

        url_button.click(
            process_url,
            inputs=[
                url_textbox,
                decode_method_radio,
                beam_size_slider,
                max_symbols_per_step_slider,
            ],
            outputs=[url_output, url_html_info],
        )
    gr.Markdown(description)


if __name__ == "__main__":
    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"

    logging.basicConfig(format=formatter, level=logging.INFO)

    demo.launch()