import os
import whisper
from moviepy.editor import VideoFileClip
from whisper.utils import get_writer

def extract_audio(video_path, audio_temp_path):
    print("Extracting uncompressed WAV audio...")
    video = VideoFileClip(video_path)
    # Extract as WAV, strictly setting the sample rate to 16000Hz for Whisper
    video.audio.write_audiofile(
        audio_temp_path, 
        codec="pcm_s16le", 
        ffmpeg_params=["-ar", "16000", "-ac", "1"]
    )
    video.close()

def transcribe_to_srt(audio_path, output_directory, output_filename):
    print("Loading Whisper model...")
    model = whisper.load_model("base")
    
    print("Transcribing audio...")
    result = model.transcribe(audio_path)
    
    print("Creating SRT file...")
    srt_writer = get_writer("srt", output_directory)
    srt_writer(result, output_filename, {})
    
    print(f"SRT subtitle file saved in: {output_directory}")

if __name__ == "__main__":
    video_path = "my_video.mp4" 
    audio_temp_path = "temp_audio.wav"  # Changed extension to .wav
    
    output_dir = os.getcwd()  
    output_prefix = "my_video_subtitles"  
    
    extract_audio(video_path, audio_temp_path)
    transcribe_to_srt(audio_temp_path, output_dir, output_prefix)
    
    if os.path.exists(audio_temp_path):
        os.remove(audio_temp_path)