import torch
from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read
import gradio as gr

MODEL_NAME = "EwoutLagendijk/whisper-small-indonesian"
BATCH_SIZE = 8

device = 0 if torch.cuda.is_available() else "cpu"

pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    chunk_length_s=30,
    device=device,
)


# Copied from https://github.com/openai/whisper/blob/c09a7ae299c4c34c5839a76380ae407e7d785914/whisper/utils.py#L50
def format_timestamp(seconds: float, always_include_hours: bool = False, decimal_marker: str = "."):
    if seconds is not None:
        milliseconds = round(seconds * 1000.0)

        hours = milliseconds // 3_600_000
        milliseconds -= hours * 3_600_000

        minutes = milliseconds // 60_000
        milliseconds -= minutes * 60_000

        seconds = milliseconds // 1_000
        milliseconds -= seconds * 1_000

        hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
        return f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
    else:
        # we have a malformed timestamp so just return it as is
        return seconds


def transcribe_speech(filepath):
    # Load the audio
    audio, sampling_rate = librosa.load(filepath, sr=16000)

    # Define chunk size (e.g., 30 seconds)
    chunk_duration = 30  # in seconds
    chunk_samples = chunk_duration * sampling_rate

    # Process audio in chunks
    transcription = []
    for i in range(0, len(audio), chunk_samples):
        chunk = audio[i:i + chunk_samples]

        # Convert the chunk into input features
        inputs = processor(audio=chunk, sampling_rate=16000, return_tensors="pt").input_features

        # Generate transcription for the chunk
        generated_ids = model.generate(
            inputs,
            max_new_tokens=444,  # Max allowed by Whisper
            forced_decoder_ids=processor.get_decoder_prompt_ids(language="id", task="transcribe")
        )

        # Decode and append the transcription
        chunk_transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        transcription.append(chunk_transcription)

    # Combine all chunk transcriptions into a single string
    return " ".join(transcription)



demo = gr.Blocks()

mic_transcribe = gr.Interface(
    fn=transcribe_speech,
    inputs=gr.Audio(sources="microphone", type="filepath"),
    outputs=gr.components.Textbox(),
)

file_transcribe = gr.Interface(
    fn=transcribe_speech,
    inputs=gr.Audio(sources="upload", type="filepath"),
    outputs=gr.components.Textbox(),
)

with demo:
    gr.TabbedInterface([mic_transcribe, file_transcribe], ["Transcribe Microphone", "Transcribe Audio File"])

demo.launch(share=True, debug=True)