import os
import cv2
import numpy as np
from flask import Flask, render_template, request, send_from_directory
from werkzeug.utils import secure_filename

# Define the folder paths for saving input and output videos
UPLOAD_FOLDER = 'input_videos/'
OUTPUT_FOLDER = 'output_videos/'

# Allowed video extensions
ALLOWED_EXTENSIONS = {'mp4', 'avi', 'mov'}

# Initialize the Flask app
app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['OUTPUT_FOLDER'] = OUTPUT_FOLDER


def allowed_file(filename):
    """Check if the uploaded file is allowed based on its extension."""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def extract_and_compare(video1_path, video2_path, output_path):
    """Compare the two videos and save the output video with differences highlighted."""
    cap1 = cv2.VideoCapture(video1_path)
    cap2 = cv2.VideoCapture(video2_path)
    
    fps = cap1.get(cv2.CAP_PROP_FPS)
    width = int(cap1.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap1.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    while cap1.isOpened() and cap2.isOpened():
        ret1, frame1 = cap1.read()
        ret2, frame2 = cap2.read()
        
        if not ret1 or not ret2:
            break
        
        frame1_gray = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
        frame2_gray = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
        
        diff = cv2.absdiff(frame1_gray, frame2_gray)
        _, diff_thresh = cv2.threshold(diff, 30, 255, cv2.THRESH_BINARY)
        
        contours, _ = cv2.findContours(diff_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        for contour in contours:
            if cv2.contourArea(contour) > 500:
                (x, y, w, h) = cv2.boundingRect(contour)
                cv2.rectangle(frame1, (x, y), (x + w, y + h), (0, 255, 0), 2)
        
        out.write(frame1)
    
    cap1.release()
    cap2.release()
    out.release()


@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        # Check if files are part of the request
        if 'video1' not in request.files or 'video2' not in request.files:
            return 'No files part'

        video1 = request.files['video1']
        video2 = request.files['video2']

        # Check if the files are allowed
        if video1 and allowed_file(video1.filename) and video2 and allowed_file(video2.filename):
            # Secure filenames and save to input_videos folder
            filename1 = secure_filename(video1.filename)
            filename2 = secure_filename(video2.filename)
            video1_path = os.path.join(app.config['UPLOAD_FOLDER'], filename1)
            video2_path = os.path.join(app.config['UPLOAD_FOLDER'], filename2)
            video1.save(video1_path)
            video2.save(video2_path)

            # Output path
            output_path = os.path.join(app.config['OUTPUT_FOLDER'], 'output_with_highlights.avi')

            # Compare the videos and generate the output
            extract_and_compare(video1_path, video2_path, output_path)

            # Return the output video
            return send_from_directory(app.config['OUTPUT_FOLDER'], 'output_with_highlights.avi', as_attachment=True)

    return render_template('index.html')


if __name__ == '__main__':
    app.run(debug=True)