import gradio as gr
import torchaudio
import torchaudio.transforms as T
from transformers import pipeline
import requests
from pydub import AudioSegment
from pydub.silence import split_on_silence
import io
import os
from bs4 import BeautifulSoup
import re
import numpy as np
from moviepy import editor
import soundfile as sf

# Load the transcription model
transcription_pipeline = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h")

def download_audio_from_url(url):
    try:
        if "share" in url:
            print("Processing shareable link...")
            response = requests.get(url)
            soup = BeautifulSoup(response.content, 'html.parser')
            video_tag = soup.find('video')
            if video_tag and 'src' in video_tag.attrs:
                video_url = video_tag['src']
                print(f"Extracted video URL: {video_url}")
            else:
                raise ValueError("Direct video URL not found in the shareable link.")
        else:
            video_url = url
        
        print(f"Downloading video from URL: {video_url}")
        response = requests.get(video_url)
        audio_bytes = response.content
        print(f"Successfully downloaded {len(audio_bytes)} bytes of data")
        return audio_bytes
    except Exception as e:
        print(f"Error in download_audio_from_url: {str(e)}")
        raise

def transcribe_audio(video_bytes):
    try:
        # Save the video bytes to a temporary file
        with open("temp_video.mp4", "wb") as f:
            f.write(video_bytes)

        # Extract audio from video
        video = editor.VideoFileClip("temp_video.mp4")
        audio = video.audio

        # Export audio as mono WAV
        audio.write_audiofile("temp_audio.wav", fps=16000, nbytes=2, codec='pcm_s16le')

        # Load the audio file
        audio_data, sample_rate = sf.read("temp_audio.wav")

        # Ensure audio is mono
        if len(audio_data.shape) > 1:
            audio_data = audio_data.mean(axis=1)

        # Normalize the audio data
        audio_data = audio_data.astype(np.float32) / np.max(np.abs(audio_data))

        # Transcribe the audio
        result = transcription_pipeline(audio_data)
        transcript = result['text']

        # Clean up temporary files
        os.remove("temp_video.mp4")
        os.remove("temp_audio.wav")

        return transcript
    except Exception as e:
        print(f"Error in transcribe_audio: {str(e)}")
        raise

def transcribe_video(url):
    try:
        print(f"Attempting to download audio from URL: {url}")
        audio_bytes = download_audio_from_url(url)
        print(f"Successfully downloaded {len(audio_bytes)} bytes of audio data")
        
        print("Starting audio transcription...")
        transcript = transcribe_audio(audio_bytes)
        print("Transcription completed successfully")
        
        return transcript
    except Exception as e:
        error_message = f"An error occurred: {str(e)}"
        print(error_message)
        return error_message
        
def download_transcript(transcript):
    return transcript, "transcript.txt"

# Create the Gradio interface
with gr.Blocks(title="Video Transcription") as demo:
    gr.Markdown("# Video Transcription")
    video_url = gr.Textbox(label="Video URL")
    transcribe_button = gr.Button("Transcribe")
    transcript_output = gr.Textbox(label="Transcript", lines=20)
    download_button = gr.Button("Download Transcript")
    download_link = gr.File(label="Download Transcript")
    
    transcribe_button.click(fn=transcribe_video, inputs=video_url, outputs=transcript_output)
    download_button.click(fn=download_transcript, inputs=transcript_output, outputs=[download_link, download_link])

demo.launch()