import streamlit as st
from stable_whisper import load_model
import requests
import os

# Variables
valid_api_token = st.secrets["API_TOKEN"]

# Upload audio file
uploaded_file = st.file_uploader("Upload Audio File", type=["mp3", "wav", "mov"])

# Free tier or API token option
use_free_tier = st.checkbox("Free Tier (Max 2 minutes)")
api_token = st.text_input("API Token (Unlimited)")

# Model selection
model_size = st.selectbox("Model Size", ("tiny", "base", "small", "medium"))

def transcribe_to_subtitle(audio_bytes, model_name):
  """Transcribe audio to subtitle using OpenAI Whisper"""
  # Load model based on selection
  model = load_model(model_name)
  
  # Check file size for free tier
  if use_free_tier and len(audio_bytes) > 2 * 60 * 1024:
      st.error("Free tier only supports audio files under 2 minutes")
      return
  
  #  Transcribe audio
  result = model.transcribe(audio_bytes)
  
  # Generate subtitle file
  subtitle_text = result.text
  with open("audio.srt", "w") as outfile:
      outfile.write(subtitle_text)
  
  # Download option
  st.success("Transcription successful! Download subtitle file?")
  if st.button("Download"):
      st.write("Downloading...")
      with open("audio.srt", "rb") as f:
          st.download_button("Download Subtitle", f, "audio.srt")
      os.remove("audio.srt")  # Remove temporary file

if uploaded_file is not None:
  audio_bytes = uploaded_file.read()
  # Check for API token if free tier is not selected
  if not use_free_tier and not api_token:
      st.error("API token required for non-free tier usage")
  else:
      transcribe_to_subtitle(audio_bytes, model_size)