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="wide")

# Header
st.title("🎥 Smart Video-to-Text Summarization App")


# Initialize session state
if "video_path" not in st.session_state:
    st.session_state.video_path = None
if "transcription" not in st.session_state:
    st.session_state.transcription = None

# Upload Video Section
st.header("Upload Your Video")
upload_option = st.radio("Choose upload method:", ["📁 Local File", "📺 YouTube URL"])

if upload_option == "📁 Local File":
    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!")

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)}")

# Display video
if st.session_state.video_path:
    st.header("Preview & Process Video")
    st.video(st.session_state.video_path)

    process_btn = st.button("🚀 Process Video", key="process-video", use_container_width=True)

    if process_btn:
        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 extracting audio: {str(e)}")
                return None

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

        audio_path = extract_audio(st.session_state.video_path)
        if audio_path:
            st.session_state.transcription = transcribe_audio(audio_path)

# Summarization & Translation
if st.session_state.transcription:
    st.header("Results")

    summarize_btn = st.button("📝 Summarize Text")
    translate_btn = st.button("🌍 Translate Summary")

    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 = None
    if summarize_btn:
        summary = summarize_text(st.session_state.transcription)

    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

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

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