import cv2
import torch
import numpy as np


def extract_video_features(video_path, model, processor, device):
    """提取视频关键帧特征

    Args:
        video_path: 视频文件路径
        model: CLIP模型
        processor: CLIP处理器
        device: 计算设备

    Returns:
        features: 形状为 (n_frames, feature_dim) 的特征数组
        timestamps: 对应帧的时间戳(秒)
    """
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"无法打开视频文件: {video_path}")

    # 获取视频属性
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # 处理无效帧率
    if fps <= 0:
        fps = 30  # 默认帧率
        print(f"警告: 视频帧率无效，使用默认值 {fps}")

    # 每3秒提取一帧
    frame_interval = max(1, int(fps * 2))  # 确保至少为1
    features = []
    timestamps = []

    for i in range(0, total_frames, frame_interval):
        start_frame = i

        cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame) 
        ret, frame = cap.read()  # 读取该位置的帧（图片）

        if not ret:
            break

        try:
            # 预处理帧
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            # 提取特征
            inputs = processor(images=frame_rgb, return_tensors="pt").to(device)
            with torch.no_grad():
                frame_features = model.get_image_features(**inputs)

            # 保存结果
            features.append(frame_features.cpu().numpy().squeeze())  # 去除batch维度，二维变一维(v_row维)  array([ 9.67003047e-01,  7.20102668e-01, -6.19670749e-03, -1.22871208e+00, ...]
            timestamps.append(start_frame / fps)

        except Exception as e:
            print(f"处理帧 {i} 时出错: {str(e)}")
            continue

    cap.release()

    # 转换为numpy数组
    if len(features) > 0:
        features = np.vstack(features)  # 形状: (n_frames, feature_dim)
        timestamps = np.array(timestamps)
    else:
        features = np.empty((0, 512))  # 空数组
        timestamps = np.array([])

    return features, timestamps, total_frames, fps