import streamlit as st
from transformers import pipeline
from pytube import YouTube
import tempfile
from moviepy.editor import VideoFileClip
from PIL import Image
import requests
from io import BytesIO

st.set_page_config(page_title="AI Media Verifier", layout="wide")

@st.cache_resource
def load_image_model():
    return pipeline("image-classification", model="google/vit-base-patch16-224")

@st.cache_resource
def load_video_model():
    return pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

def detect_image(image):
    model = load_image_model()
    return model(image)

def detect_video(video_url):
    yt = YouTube(video_url)
    stream = yt.streams.filter(file_extension='mp4', progressive=True).order_by('resolution').desc().first()
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
        stream.download(filename=tmp.name)
        clip = VideoFileClip(tmp.name).subclip(0, min(10, yt.length))
        frame = clip.get_frame(clip.duration / 2)
        image = Image.fromarray(frame)
        return detect_image(image)

st.title("AI Media Verifier: Image & Video Authenticity")

option = st.radio("Choose input type:", ("Image", "Video"))

if option == "Image":
    uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
    if uploaded_image:
        image = Image.open(uploaded_image)
        st.image(image, caption="Uploaded Image", use_container_width=True)
        with st.spinner("Analyzing image..."):
            results = detect_image(uploaded_image)
        st.success("Analysis Complete")
        for res in results:
            st.write(f"**{res['label']}**: {res['score']*100:.2f}%")

elif option == "Video":
    video_url = st.text_input("Enter YouTube Video Link")
    if video_url:
        st.video(video_url)
        with st.spinner("Analyzing video..."):
            results = detect_video(video_url)
        st.success("Video Analysis Complete")
        for res in results:
            st.write(f"**{res['label']}**: {res['score']*100:.2f}%")