import os
import streamlit as st
from moviepy.video.io.VideoFileClip import VideoFileClip
from pydub import AudioSegment
import whisper
from transformers import pipeline, MarianMTModel, MarianTokenizer
import yt_dlp as youtube_dl


# App Configuration
st.set_page_config(page_title="Video-to-Text Summarization", layout="centered")

# Header
st.title("🎥 Smart Video-to-Text Summarization App")
st.markdown("""
This app helps you:
- Convert videos into text and summarize them.
- Extract multilingual transcriptions and translations.
- Process videos with multiple speakers.
""")

# Temporary video storage
if "video_path" not in st.session_state:
    st.session_state.video_path = None

# 1. Upload Video Section
st.header("Upload Your Video")

# Choose upload option
upload_option = st.selectbox("Select Upload Method", ["Local", "YouTube URL"])

# Upload Local File
if upload_option == "Local":
    video_file = st.file_uploader("Upload your video file", type=["mp4", "mkv", "avi"])
    if video_file:
        with open("uploaded_video.mp4", "wb") as f:
            f.write(video_file.read())
        st.session_state.video_path = "uploaded_video.mp4"
        st.success("Video uploaded successfully!")

# Download Video from YouTube
elif upload_option == "YouTube URL":
    youtube_url = st.text_input("Enter YouTube URL")
    if youtube_url:
        try:
            os.system(f"yt-dlp -o video.mp4 {youtube_url}")
            st.session_state.video_path = "video.mp4"
            st.success("YouTube video downloaded successfully!")
        except Exception as e:
            st.error(f"Error downloading video: {str(e)}")

# 2. Process Video Section (After Upload)
if st.session_state.video_path:
    st.header("Process Your Video")
    st.write(f"Processing {st.session_state.video_path}...")

    # Extract Audio from Video
    def extract_audio(video_path):
        try:
            audio = AudioSegment.from_file(video_path)
            audio.export("extracted_audio.mp3", format="mp3")
            st.success("Audio extracted successfully!")
            return "extracted_audio.mp3"
        except Exception as e:
            st.error(f"Error in extracting audio: {str(e)}")
            return None

    audio_path = extract_audio(st.session_state.video_path)

    # Real-time Audio Transcription
    def transcribe_audio(audio_path):
        try:
            model = whisper.load_model("base")
            result = model.transcribe(audio_path)
            st.text_area("Transcription", result['text'], height=200)
            return result['text']
        except Exception as e:
            st.error(f"Error in transcription: {str(e)}")
            return None

    if audio_path:
        transcription = transcribe_audio(audio_path)

# 3. Summarize and Translate
if 'transcription' in locals():
    st.header("Results")

    # Summarize Text
    def summarize_text(text):
        try:
            summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
            summary = summarizer(text, max_length=150, min_length=30, do_sample=False)
            st.text_area("Summary", summary[0]['summary_text'], height=150)
            return summary[0]['summary_text']
        except Exception as e:
            st.error(f"Error in summarization: {str(e)}")
            return None

    summary = summarize_text(transcription)

    # Translate Text
    def translate_text(text, src_lang="en", tgt_lang="es"):
        try:
            model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
            tokenizer = MarianTokenizer.from_pretrained(model_name)
            model = MarianMTModel.from_pretrained(model_name)
            translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True))
            translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
            st.text_area("Translated Summary", translated_text, height=150)
            return translated_text
        except Exception as e:
            st.error(f"Error in translation: {str(e)}")
            return None

    target_language = st.selectbox("Select Translation Language", ["es", "fr", "de", "zh"])
    if target_language:
        translated_summary = translate_text(summary, tgt_lang=target_language)

else:
    st.info("Please upload a video to start the process.")