import tempfile
from pathlib import Path

import gradio as gr
from faster_whisper import WhisperModel

MODEL_SIZE = "base"

model = WhisperModel(
    MODEL_SIZE,
    device="cpu",
    compute_type="int8"
)

LANGUAGE_NAMES = {
    "fr": "French",
    "en": "English",
    "de": "German",
    "fa": "Persian",
    "es": "Spanish",
    "it": "Italian",
    "pt": "Portuguese",
    "nl": "Dutch"
}


def transcribe(audio_file):

    if audio_file is None:
        return "", "", None

    segments, info = model.transcribe(
        audio_file,
        beam_size=5
    )

    transcript_lines = []
    timestamp_lines = []

    for segment in segments:
        transcript_lines.append(segment.text)

        timestamp_lines.append(
            f"[{segment.start:.2f}s ? {segment.end:.2f}s] "
            f"{segment.text}"
        )

    transcript = "\n".join(transcript_lines)

    transcript_with_timestamps = "\n".join(
        timestamp_lines
    )

    detected_language = LANGUAGE_NAMES.get(
        info.language,
        info.language
    )

    summary = (
        f"Detected language: {detected_language}\n"
        f"Confidence: {info.language_probability:.2%}"
    )

    output_file = Path(tempfile.gettempdir()) / "transcript.txt"

    with open(
        output_file,
        "w",
        encoding="utf-8"
    ) as f:
        f.write(transcript_with_timestamps)

    return (
        summary,
        transcript,
        str(output_file)
    )


with gr.Blocks(title="EchoScript") as demo:

    gr.Markdown(
        """
        # EchoScript

        Upload an audio file and automatically
        transcribe speech to text.
        """
    )

    audio_input = gr.Audio(
        type="filepath",
        label="Upload Audio"
    )

    transcribe_button = gr.Button(
        "Transcribe"
    )

    language_output = gr.Textbox(
        label="Language Information"
    )

    transcript_output = gr.Textbox(
        label="Transcript",
        lines=20
    )

    download_output = gr.File(
        label="Download Transcript"
    )

    transcribe_button.click(
        fn=transcribe,
        inputs=audio_input,
        outputs=[
            language_output,
            transcript_output,
            download_output
        ]
    )

demo.launch()