import spaces
import os
import tempfile
import numpy as np
import torch
import librosa
import soundfile as sf
import gradio as gr
from transformers import AutoProcessor, AutoModelForCTC

CHUNK_SEC   = 20
OVERLAP_SEC = 1.0
MODEL_ID    = os.environ.get("MODEL_ID")

processor = AutoProcessor.from_pretrained(MODEL_ID)
model     = AutoModelForCTC.from_pretrained(MODEL_ID).eval()


def convert_to_wav(input_path):
    out = tempfile.mktemp(suffix=".wav")
    audio, _ = librosa.load(input_path, sr=16000, mono=True)
    sf.write(out, audio, 16000, subtype="PCM_16")
    return out


@spaces.GPU
def transcribe(audio_path, progress=gr.Progress()):
    model.to("cuda")

    wav       = convert_to_wav(audio_path)
    audio, sr = librosa.load(wav, sr=16000, mono=True)

    chunk_samples   = int(CHUNK_SEC   * sr)
    overlap_samples = int(OVERLAP_SEC * sr)
    step_samples    = chunk_samples - overlap_samples
    starts          = list(range(0, len(audio), step_samples))

    parts = []

    for i, start in enumerate(starts):
        progress(i / len(starts), desc=f"Transcribing chunk {i+1}/{len(starts)}")
        end   = min(start + chunk_samples, len(audio))
        chunk = audio[start:end]

        if np.sqrt(np.mean(chunk**2)) < 0.001:
            continue

        inputs      = processor(chunk, sampling_rate=16000, return_tensors="pt")
        input_feats = inputs.input_features.to("cuda")

        with torch.no_grad():
            logits = model(input_feats).logits

        text = processor.batch_decode(torch.argmax(logits, dim=-1))[0].strip()
        parts.append(text)

        del logits, input_feats, inputs
        torch.cuda.empty_cache()

    transcript = " ".join(parts).strip()

    txt = tempfile.mktemp(suffix=".txt")
    with open(txt, "w", encoding="utf-8") as f:
        f.write(transcript + "\n")

    return transcript, gr.update(value=txt, visible=True)


with gr.Blocks(title="Chichewa ASR") as demo:
    gr.Markdown("## Chichewa Speech Transcription")

    audio_input    = gr.Audio(sources=["upload", "microphone"], type="filepath")
    transcribe_btn = gr.Button("Transcribe")
    transcript_out = gr.Textbox(label="Transcription", lines=6, interactive=False)
    download_out   = gr.File(label="Download transcript", visible=False)

    transcribe_btn.click(
        fn=transcribe,
        inputs=audio_input,
        outputs=[transcript_out, download_out],
        show_progress="full",
    )

demo.launch(debug=True)