import cv2
import numpy as np
import time
import random
from PIL import Image
import argparse
import subprocess
from transparent_background import Remover
import os
import torch


def process_video(input_video, output_video, mode='Normal'):
    if mode == 'Fast':
        remover = Remover(mode='fast')
    else:
        remover = Remover()

    cap = cv2.VideoCapture(input_video)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # Get total frames
    processed_frames = 0
    start_time = time.time()

    # Get video properties
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Start ffmpeg subprocess
    ffmpeg_command = [
        'ffmpeg',
        '-y',  # Overwrite output file if it exists
        '-f', 'rawvideo',
        '-vcodec', 'rawvideo',
        '-pix_fmt', 'rgb24',
        '-s', f'{width}x{height}',  # Size of one frame
        '-r', str(fps),  # Frames per second
        '-i', '-',  # Input from pipe
        '-c:v', 'libx264',
        '-crf', '0',
        output_video
    ]

    proc = subprocess.Popen(ffmpeg_command, stdin=subprocess.PIPE)

    while cap.isOpened():
        ret, frame = cap.read()

        if not ret:
            break

        if time.time() - start_time >= 20 * 60 - 5:
            print("GPU Timeout is coming")
            cap.release()
            proc.stdin.close()
            proc.wait()
            return output_video

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(frame).convert('RGB')

        processed_frames += 1
        print(f"Processing frame {processed_frames}/{total_frames}")
        out = remover.process(img, type='green')
        proc.stdin.write(np.array(out).tobytes())

    cap.release()
    proc.stdin.close()
    proc.wait()

    print(f"Output video saved to {output_video}")
    return output_video

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Remove background from video using transparent_background library.")
    parser.add_argument('input', type=str, help='Input video file path')
    parser.add_argument('output', type=str, help='Output video file path')
    parser.add_argument('--mode', type=str, default='Normal', choices=['Fast', 'Normal'], help='Mode of operation')
    args = parser.parse_args()

    process_video(args.input, args.output, args.mode)