import os
import subprocess

import cv2

import functools
import time


def timer_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        print(f"{func.__name__} took {execution_time:.2f} seconds to execute")
        return result

    return wrapper


@timer_decorator
def extract_frames_by_rate(video_path, output_folder, frame_rate):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    cmd = [
        'ffmpeg',
        '-i', video_path,
        '-vf', f'fps={frame_rate}',
        os.path.join(output_folder, 'frame_%05d.jpg')
    ]
    subprocess.run(cmd, check=True)

# def extract_frames_by_rate(video_path, output_folder, frame_rate):
#     """
#     Extracts frames from a video at a specified frame rate.
#
#     Args:
#         video_path (str): Path to the input video file.
#         output_folder (str): Directory to save the extracted frames.
#         frame_rate (int): Number of frames to extract per second of the video.
#     """
#     # Ensure the output directory exists
#     if not os.path.exists(output_folder):
#         os.makedirs(output_folder)
#
#     # Load the video
#     video = cv2.VideoCapture(video_path)
#
#     # Check if the video is opened successfully
#     if not video.isOpened():
#         print(f"Error: Cannot open video file {video_path}")
#         return
#
#     # Get video properties
#     fps = int(video.get(cv2.CAP_PROP_FPS))  # Frames per second
#     total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))  # Total number of frames
#     duration = total_frames / fps  # Duration in seconds
#
#     print(f"Video loaded: {video_path}")
#     print(f"Total Frames: {total_frames}, FPS: {fps}, Duration: {duration:.2f} seconds")
#
#     # Calculate frame interval (in terms of frame number)
#     frame_interval = fps // frame_rate
#
#     # Frame counter
#     frame_count = 0
#     saved_count = 0
#
#     while True:
#         # Read a frame
#         ret, frame = video.read()
#
#         # Break the loop if the video ends
#         if not ret:
#             break
#
#         # Save frame if it matches the frame interval
#         if frame_count % frame_interval == 0:
#             frame_filename = os.path.join(output_folder, f"frame_{saved_count:05d}.jpg")
#             cv2.imwrite(frame_filename, frame)
#             saved_count += 1
#
#         frame_count += 1
#
#     # Release video resources
#     video.release()
#     print(f"Extraction complete. Total frames saved: {saved_count}. FPS used to extracted: {frame_rate}")


@timer_decorator
def create_video_from_frames(input_folder, output_video_path, frame_rate, resolution=(360, 640)):
    """
       Creates a video from preprocessed frames.

       Args:
           input_folder (str): Path to the folder containing frames.
           output_video_path (str): Path to save the output video.
           frame_rate (int): Frames per second for the output video.
           resolution (tuple): Resolution of the output video (width, height).
       """
    # Get sorted list of image files in the folder
    frame_files = sorted(
        [f for f in os.listdir(input_folder) if f.lower().endswith(('.jpg', '.png'))]
    )

    # Initialize the video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_writer = cv2.VideoWriter(output_video_path, fourcc, frame_rate, resolution)

    for frame_file in frame_files:
        frame_path = os.path.join(input_folder, frame_file)
        frame = cv2.imread(frame_path)

        if frame is None:
            print(f"Error reading frame: {frame_path}")
            continue

        # Get the frame's original resolution
        original_height, original_width = frame.shape[:2]

        # **Check if resizing is needed**
        if (original_width, original_height) != resolution:
            # Ensure the frame matches the target resolution
            frame = cv2.resize(frame, resolution, interpolation=cv2.INTER_CUBIC)

        # Write the frame to the video
        video_writer.write(frame)

    # Release the video writer
    video_writer.release()
    print(f"Video saved to: {output_video_path}")