import cv2
import numpy as np
import subprocess
import json
import time


def get_video_dimensions(video_path):
    """ 获取视频的宽高信息 """
    command = [
        'ffprobe',
        '-v', 'error',
        '-select_streams', 'v:0',
        '-show_entries', 'stream=width,height',
        '-of', 'json',
        video_path
    ]
    result = subprocess.run(command, capture_output=True, text=True)
    info = json.loads(result.stdout)
    width = info['streams'][0]['width']
    height = info['streams'][0]['height']
    return width, height


def get_video_fps(video_path):
    command = [
        'ffprobe',
        '-v', 'error',
        '-select_streams', 'v:0',
        '-show_entries', 'stream=r_frame_rate',
        '-of', 'json',
        video_path
    ]
    result = subprocess.run(command, capture_output=True, text=True)
    info = json.loads(result.stdout)
    rate = info['streams'][0]['r_frame_rate']
    num, denom = map(int, rate.split('/'))
    return num / denom


def decode_and_display_video_with_ffmpeg(video_path):
    width, height = get_video_dimensions(video_path)
    frame_size = width * height * 3
    fps = get_video_fps(video_path)
    frame_delay = 1.0 / fps

    command = [
        'ffmpeg',
        '-hwaccel', 'cuda',
        '-i', video_path,
        '-f', 'rawvideo',
        '-pix_fmt', 'bgr24',
        '-vsync', '2',
        '-'
    ]

    print("Starting FFmpeg process...")
    video_process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=10 ** 8)

    try:
        last_frame_time = time.time()
        while True:
            start_time = time.time()

            in_bytes = video_process.stdout.read(frame_size)
            if not in_bytes:
                print("Video ended.")
                break

            frame = np.frombuffer(in_bytes, np.uint8).reshape((height, width, 3))
            cv2.imshow('Video', frame)

            current_time = time.time()
            elapsed_time = current_time - last_frame_time
            delay = frame_delay - elapsed_time

            if delay > 0:
                time.sleep(delay)

            processing_time = time.time() - start_time
            print(f"Processing time for the frame: {processing_time:.4f} seconds")

            last_frame_time = current_time

            # Check for 'q' key press
            key = cv2.waitKey(1)
            if key & 0xFF == ord('q'):
                print("Exiting on user request.")
                break

    except Exception as e:
        print(f"An error occurred: {e}")

    finally:
        video_process.stdout.close()
        video_process.stderr.close()
        video_process.terminate()
        video_process.wait()
        cv2.destroyAllWindows()


video_path = 'E:/BaiduNetdiskDownload/out.avi'
decode_and_display_video_with_ffmpeg(video_path)