import cv2
import numpy as np
import os
from rembg import remove
from PIL import Image

def process_video(input_path, output_path):
    # Check if input exists
    if not os.path.exists(input_path):
        print(f"Error: Input file '{input_path}' not found.")
        return

    print(f"Processing video: {input_path}")
    
    cap = cv2.VideoCapture(input_path)
    
    # Get video properties
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    print(f"Resolution: {width}x{height}, FPS: {fps}, Total Frames: {total_frames}")

    # Initialize video writer
    # Using 'mp4v' codec for MP4. Note that standard MP4 does not support alpha channel easily.
    # For web transparency, we usually need WebM with VP9 or a specific MOV codec (ProRes 4444).
    # Here we will try to create a WebM file (VP9) which supports alpha.
    
    fourcc = cv2.VideoWriter_fourcc(*'VP90') 
    output_ext = os.path.splitext(output_path)[1].lower()
    
    if output_ext == '.mp4':
        print("Warning: MP4 container often doesn't support alpha transparency widely. Switching codec might be needed.")
        # Try mp4v just in case, but alpha might be lost or black
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    elif output_ext == '.webm':
         fourcc = cv2.VideoWriter_fourcc(*'VP90')

    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    frame_count = 0
    
    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Convert BGR (OpenCV) to RGB (PIL/rembg)
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            pil_im = Image.fromarray(frame_rgb)
            
            # Remove background using rembg
            output_pil = remove(pil_im)
            
            # Convert back to numpy
            output_np = np.array(output_pil)
            
            # Convert RGB to BGR for OpenCV handling (if we were saving a normal video)
            # But wait, OpenCV VideoWriter expects BGR. 
            # If we want transparency, we need a 4-channel write.
            # Standard cv2.VideoWriter might struggle with 4 channels depending on backend.
            
            # Let's check if the output has alpha
            if output_np.shape[2] == 4:
                # If we are writing to a format that supports alpha (like VP9 WebM), we should pass the alpha.
                # However, basic cv2 VideoWriter might not support 4 channels.
                # A safer bet for a simple script is to save as a sequence of PNGs or find a writer that supports it.
                # For this PoC, let's try writing the frame. 
                
                # If VideoWriter fails with 4 channels, we fallback to black background.
                frame_bgr_alpha = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA)
                out.write(frame_bgr_alpha)
            else:
                frame_bgr = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGR)
                out.write(frame_bgr)
            
            frame_count += 1
            if frame_count % 10 == 0:
                print(f"Processed {frame_count}/{total_frames} frames...")

    except Exception as e:
        print(f"Error processing frames: {e}")

    finally:
        cap.release()
        out.release()
        print("Done.")

if __name__ == "__main__":
    # Ensure dependencies: pip install rembg opencv-python pillow numpy
    input_file = "cup_spin.mp4"
    output_file = "cup_spin_no_bg.webm" # using webm for transparency support
    
    process_video(input_file, output_file)