import streamlit as st
import tempfile
import os
import cv2
import numpy as np

def remove_watermark(frame):
    # Convert to grayscale
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    # Thresholding to create a mask
    _, mask = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
    # Inpaint to remove watermark
    result = cv2.inpaint(frame, mask, 3, cv2.INPAINT_TELEA)
    return result

def enhance_resolution(input_path, output_path):
    os.system(f"ffmpeg -i {input_path} -vf scale=1920:1080 {output_path}")

def trim_invidea_ai_branding(input_path, output_path, duration=None):
    if duration is None:
        duration = get_video_duration(input_path) - 3  # Trim last 3 seconds
    os.system(f"ffmpeg -i {input_path} -t {duration} {output_path}")

def get_video_duration(video_path):
    import ffmpeg
    probe = ffmpeg.probe(video_path)
    return float(probe['format']['duration'])

def process_video(input_video):
    cap = cv2.VideoCapture(input_video)
    if not cap.isOpened():
        return None

    temp_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(temp_output.name, fourcc, 30.0, (int(cap.get(3)), int(cap.get(4))))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        processed_frame = remove_watermark(frame)
        out.write(processed_frame)

    cap.release()
    out.release()

    # Enhance resolution
    enhanced_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    enhance_resolution(temp_output.name, enhanced_output.name)

    # Trim branding
    final_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    trim_invidea_ai_branding(enhanced_output.name, final_output.name)

    return final_output.name

st.title("AI Video Enhancement App 🎥")

uploaded_file = st.file_uploader("Upload a video", type=["mp4", "avi", "mov"])

if uploaded_file:
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file:
        temp_file.write(uploaded_file.read())
        temp_file_path = temp_file.name

    st.write("Processing video...")
    processed_video_path = process_video(temp_file_path)

    if processed_video_path:
        st.video(processed_video_path)
        with open(processed_video_path, "rb") as file:
            st.download_button("Download Processed Video", file, "enhanced_video.mp4", "video/mp4")