streamlit_ui.py

import os
import tempfile
import subprocess
import re
from functools import lru_cache
from typing import Any

import streamlit as st
from transformers import pipeline
import requests

# Page Config
st.set_page_config(
    page_title="Swecha Telugu Audio Tracker",
    page_icon="🎙️",
    layout="centered"
)

# Constants & Env
MODEL_ID = os.getenv("MODEL_ID", "viswamaicoe/swecha-gonthuka-asr")
ASR_DEVICE = os.getenv("ASR_DEVICE", "cpu").lower()
SWECHA_API_BASE = os.getenv("SWECHA_API_BASE", "https://api.corpus.swecha.org")
SWECHA_UPLOAD_PATH = os.getenv("SWECHA_UPLOAD_PATH", "/api/v1/content")
SWECHA_AUTH_TOKEN = os.getenv("SWECHA_AUTH_TOKEN", "")

# --- Logic from backend/main.py ---

@lru_cache(maxsize=1)
def get_asr_pipeline():
    device = 0 if ASR_DEVICE == "cuda" else -1
    return pipeline(
        task="automatic-speech-recognition",
        model=MODEL_ID,
        device=device,
        chunk_length_s=30,
        batch_size=8,
    )

def clean_noisy_telugu(text: str) -> str:
    if not text:
        return ""
    cleaned = re.sub(r'్{2,}', '్', text)
    cleaned = re.sub(r'\s+', ' ', cleaned).strip()
    return cleaned

def transcribe_audio(raw_bytes: bytes, suffix: str) -> str:
    asr = get_asr_pipeline()
    
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as input_file:
        input_file.write(raw_bytes)
        input_path = input_file.name

    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as output_file:
        output_path = output_file.name

    try:
        ffmpeg_command = [
            "ffmpeg", "-y", "-i", input_path,
            "-acodec", "pcm_s16le", "-ac", "1", "-ar", "16000",
            output_path,
        ]
        subprocess.run(ffmpeg_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        result = asr(output_path, generate_kwargs={"task": "transcribe", "language": "telugu"})
        
        text = ""
        if isinstance(result, dict) and "text" in result:
            text = clean_noisy_telugu(result["text"])
        elif isinstance(result, str):
            text = clean_noisy_telugu(result)
        return text

    finally:
        for path in (input_path, output_path):
            if os.path.exists(path):
                os.remove(path)

def push_to_swecha(audio_bytes, filename, transcript, title, description):
    if not SWECHA_AUTH_TOKEN:
        return {"error": "SWECHA_AUTH_TOKEN is not configured"}

    url = f"{SWECHA_API_BASE.rstrip('/')}/{SWECHA_UPLOAD_PATH.lstrip('/')}"
    headers = {"Authorization": f"Bearer {SWECHA_AUTH_TOKEN}"}
    files = {"audio": (filename, audio_bytes, "audio/webm")}
    data = {"title": title, "description": description, "transcript": transcript}
    
    resp = requests.post(url, headers=headers, files=files, data=data, timeout=60)
    return resp.json() if resp.ok else {"error": resp.text}

# --- Streamlit UI ---

st.title("🎙️ Swecha Telugu Audio Tracker")
st.markdown("Convert Telugu audio to text and store it in the Swecha Corpus.")

tab1, tab2 = st.tabs(["Upload/Record", "Settings"])

with tab2:
    st.header("Configuration")
    model_id = st.text_input("ASR Model ID", MODEL_ID)
    auth_token = st.text_input("Swecha Auth Token", SWECHA_AUTH_TOKEN, type="password")
    if st.button("Save Settings"):
         os.environ["MODEL_ID"] = model_id
         os.environ["SWECHA_AUTH_TOKEN"] = auth_token
         st.success("Settings updated for this session!")

with tab1:
    audio_file = st.file_uploader("Choose an audio file", type=["wav", "mp3", "webm", "m4a"])
    
    # Simple Record placeholder since custom components like streamlit-mic-recorder 
    # might need specific installation and configuration.
    st.info("You can also record audio if you have 'streamlit-mic-recorder' installed. For now, please upload a file.")

    if audio_file:
        st.audio(audio_file)
        
        with st.expander("Metadata (Optional for storage)"):
            title = st.text_input("Title", value=audio_file.name)
            desc = st.text_area("Description")

        if st.button("Transcribe", type="primary"):
            with st.spinner("Transcribing Telugu..."):
                try:
                    text = transcribe_audio(audio_file.read(), os.path.splitext(audio_file.name)[1])
                    st.subheader("Transcription:")
                    st.write(text)
                    
                    if text and SWECHA_AUTH_TOKEN:
                        if st.button("Push to Swecha Corpus"):
                            res = push_to_swecha(audio_file.getvalue(), audio_file.name, text, title, desc)
                            st.json(res)
                except Exception as e:
                    st.error(f"Error: {e}")