import streamlit as st
import cv2
import numpy as np
from tensorflow.keras.models import load_model
from PIL import Image
import io
import os
from testing2 import combined_prediction, predict_video

# Load the model
model = load_model("deepfake_detector.h5")

st.set_page_config(page_title="Deepfake Detection System", layout="wide")

def process_image_in_memory(file):
    image = Image.open(file).convert("RGB")
    # Convert PIL Image to bytes
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='JPEG')
    img_byte_arr = img_byte_arr.getvalue()
    
    # Create temporary file, process it, and delete immediately
    temp_path = "temp_image.jpg"
    with open(temp_path, "wb") as f:
        f.write(img_byte_arr)
    
    results = combined_prediction(temp_path)
    
    # Clean up
    if os.path.exists(temp_path):
        os.remove(temp_path)
        
    return results, image

def process_video_in_memory(file):
    temp_path = "temp_video.mp4"
    
    # Save temporarily for processing
    with open(temp_path, "wb") as f:
        f.write(file.read())
    
    results = predict_video(temp_path)
    
    # Clean up
    if os.path.exists(temp_path):
        os.remove(temp_path)
        
    return results

def display_detailed_analysis(results):
    st.subheader("🔍 Detailed Analysis")
    
    # Main Metrics
    col1, col2 = st.columns(2)
    with col1:
        st.metric("Final Verdict", results.get("Final Prediction", "Not Available"))
        st.metric("CNN Analysis", results.get("CNN Prediction", "Not Available"))
    with col2:
        st.metric("Confidence Score", f"{results.get('Confidence Score', 0)*100:.1f}%")
        st.metric("Metadata Check", results.get("Metadata Analysis", "Not Available"))
    
    # Technical Analysis
    st.write("📊 Technical Analysis:")
    st.write(f"- **Artifact Detection**: {results.get('Artifact Analysis', 'Not Available')}")
    st.write(f"- **Noise Pattern Analysis**: {results.get('Noise Pattern Analysis', 'Not Available')}")
    
    # Symmetry Analysis if available
    if "Symmetry Analysis" in results:
        st.write("🎯 Symmetry Measurements:")
        symmetry = results["Symmetry Analysis"]
        st.write(f"- **Vertical Symmetry**: {symmetry.get('Vertical Symmetry', 0)*100:.1f}%")
        st.write(f"- **Horizontal Symmetry**: {symmetry.get('Horizontal Symmetry', 0)*100:.1f}%")

def main():
    st.title("Deepfake Detection System")
    st.write("Upload an image or video to detect if it's real or manipulated")

    file = st.file_uploader("Choose a file", type=['jpg', 'jpeg', 'png', 'mp4', 'avi'])
    
    if file:
        file_type = file.type.split('/')[0]
        
        if file_type == 'image':
            results, image = process_image_in_memory(file)
            st.image(image, caption="Uploaded Image", width=500)
            
            if st.button("Analyze Image"):
                with st.spinner("Analyzing..."):
                    display_detailed_analysis(results)
                    
        elif file_type == 'video':
            col1, col2, col3 = st.columns([1, 2, 1])
            with col2:
                # Set a smaller width for the video container
                st.markdown('<div style="width: 300px; margin: auto;">', unsafe_allow_html=True)
                st.video(file, format="video/mp4", start_time=0)
                st.markdown('</div>', unsafe_allow_html=True)
            
            if st.button("Analyze Video"):
                with st.spinner("Analyzing video frames..."):
                    results = process_video_in_memory(file)
                    
                    # Display results
                    st.subheader("Video Analysis Results")
                    col1, col2 = st.columns(2)
                    
                    with col1:
                        st.metric("Final Prediction", results["Final Video Prediction"])
                        st.metric("Confidence Score", f"{results['Confidence Score']*100:.2f}%")
                    
                    with col2:
                        st.metric("Fake Frames", results["Fake Frames"])
                        st.metric("Real Frames", results["Real Frames"])

                    # Detailed frame analysis
                    st.write("📊 Frame Analysis:")
                    st.write(f"- **Total Frames Analyzed**: {results['Total Frames Analyzed']}")
                    st.write(f"- **Fake Frames**: {results['Fake Frames']}")
                    st.write(f"- **Real Frames**: {results['Real Frames']}")
                    st.write(f"- **Fake Content Percentage**: {results['Fake Percentage']}%")

if __name__ == "__main__":
    main()