import streamlit as st
from transformers import pipeline
from PIL import Image
import requests
from urllib.parse import urlparse, parse_qs
from io import BytesIO

# Initialize the deepfake detection model
@st.cache_resource
def load_model():
    return pipeline("image-classification", model="Wvolf/ViT_Deepfake_Detection")

model = load_model()

def get_thumbnail_url(video_url):
    """
    Extracts the YouTube video ID and returns the thumbnail URL.
    """
    parsed_url = urlparse(video_url)
    video_id = None
    if 'youtube' in parsed_url.netloc:
        query_params = parse_qs(parsed_url.query)
        video_id = query_params.get('v', [None])[0]
    elif 'youtu.be' in parsed_url.netloc:
        video_id = parsed_url.path.lstrip('/')
    
    if video_id:
        return f"https://img.youtube.com/vi/{video_id}/maxresdefault.jpg"
    return None

def analyze_thumbnail(thumbnail_url):
    """
    Downloads the thumbnail image and analyzes it using the deepfake detection model.
    """
    response = requests.get(thumbnail_url)
    if response.status_code == 200:
        image = Image.open(BytesIO(response.content)).convert("RGB")
        results = model(image)
        return results, image
    else:
        st.error("Failed to retrieve the thumbnail image.")
        return None, None

# Streamlit UI
st.title("Deepfake Detection from YouTube Thumbnails")
video_url = st.text_input("Enter YouTube Video URL:")
if st.button("Analyze"):
    if video_url:
        thumbnail_url = get_thumbnail_url(video_url)
        if thumbnail_url:
            results, image = analyze_thumbnail(thumbnail_url)
            if results and image:
                st.image(image, caption="YouTube Video Thumbnail", use_column_width=True)
                st.subheader("Detection Results:")
                for result in results:
                    label = result['label']
                    confidence = result['score'] * 100
                    st.write(f"**{label}**: {confidence:.2f}%")
            else:
                st.error("Could not analyze the thumbnail.")
        else:
            st.error("Invalid YouTube URL. Please enter a valid URL.")
    else:
        st.warning("Please enter a YouTube video URL.")