# app.py

import streamlit as st
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
import os

# --- 1. تنظیمات اولیه و عنوان صفحه ---
st.set_page_config(
    page_title="Persian Whisper ASR",
    page_icon="🇮🇷🎙️",
    layout="centered"
)

st.title("🇮🇷 اپلیکیشن تبدیل گفتار به نوشتار فارسی (Whisper)")

st.markdown("""
این یک نسخه نمایشی برای مدل **`vhdm/whisper-large-fa-v1`** است.
فایل صوتی خود را آپلود کنید تا متن آن را مشاهده نمایید.
""")


# --- 2. بارگذاری مدل (با کش کردن برای سرعت بیشتر) ---
# این دکوراتور به Streamlit می‌گوید که مدل را فقط یک بار بارگذاری کند.
@st.cache_resource
def load_model():
    """Loads and caches the Whisper model and processor."""
    print("--- Loading model and processor for the first time ---")
    device = "cpu"
    torch_dtype = torch.float32

    model_id = "vhdm/whisper-large-fa-v1"

    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
    )
    processor = AutoProcessor.from_pretrained(model_id)

    pipe = pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        max_new_tokens=128,
        torch_dtype=torch_dtype,
        device=device,
    )
    print("--- Model loaded successfully ---")
    return pipe

# مدل را بارگذاری می‌کنیم
transcription_pipe = load_model()


# --- 3. بخش آپلود فایل و پردازش ---

st.header("فایل صوتی خود را آپلود کنید")
uploaded_file = st.file_uploader(
    "یک فایل صوتی انتخاب کنید (WAV, MP3, FLAC)...",
    type=["wav", "mp3", "m4a", "flac"]
)

if uploaded_file is not None:
    # نمایش فایل صوتی
    st.audio(uploaded_file, format='audio/wav')

    # دکمه برای شروع پردازش
    if st.button("شروع رونویسی"):
        # برای پردازش، فایل را به صورت موقت ذخیره می‌کنیم
        temp_file_path = f"./temp_{uploaded_file.name}"
        with open(temp_file_path, "wb") as f:
            f.write(uploaded_file.getbuffer())

        # نمایش پیام در حال پردازش
        with st.spinner("در حال پردازش فایل صوتی... لطفاً صبر کنید."):
            result = transcription_pipe(temp_file_path)
            transcription = result["text"]

        # نمایش نتیجه
        st.success("پردازش با موفقیت انجام شد!")
        st.subheader("متن رونویسی شده:")
        st.write(transcription)

        # حذف فایل موقت
        os.remove(temp_file_path)