import os
from pathlib import Path
import numpy as np
import supervision as sv
import cv2
from concurrent.futures import ThreadPoolExecutor, as_completed

import os
import cv2
from pathlib import Path

def generate_video_to_objects(
    obj_dict: dict[dict],
    input_video_path: str,
    output_dir: str,
) -> None:
    """
    根据 obj_dict 中的物体信息，从原视频中截取附近帧，生成新视频

    :param obj_dict: 包含物体信息的字典，每个元素结构为:
        "0": {
            "class_id": 4,
            "start_frame": 551,
            "end_frame": 597
        },
    :param input_video_path: 输入视频文件路径
    :param output_dir: 输出视频目录
    """
    # 最低秒数
    min_seconds = 2
    # 前后额外帧数
    extra_frames = 5

    print(f"开始抽取物体视频: {input_video_path}")
    print(f"输出目录: {output_dir}")
    # 确保输出目录存在
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    # 1. 打开原始视频
    backends = [
        (cv2.CAP_FFMPEG, 'FFmpeg'),
        (cv2.CAP_DSHOW, 'DirectShow'),
        (cv2.CAP_ANY, 'Default')
    ]

    cap = None
    for backend, backend_name in backends:
        try:
            cap = cv2.VideoCapture(input_video_path, backend)
            if cap.isOpened():
                if backend == cv2.CAP_FFMPEG:
                    try:
                        cap.set(cv2.CAP_PROP_HW_ACCELERATION, cv2.VIDEO_ACCELERATION_ANY)
                        hw_accel = cap.get(cv2.CAP_PROP_HW_ACCELERATION)
                        print(f"FFmpeg硬件加速: {'已启用' if hw_accel > 0 else '未启用'}")
                    except Exception as e:
                        print(f"设置硬件加速失败: {e}")
                print(f"使用后端: {backend_name}")
                break
        except Exception as e:
            print(f"尝试{backend_name}后端失败: {e}")
            continue

    if not cap or not cap.isOpened():
        raise Exception(f"无法打开视频文件: {input_video_path}")

    # 2. 获取视频参数
    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))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    min_frames = int(min_seconds * fps)

    fourcc = cv2.VideoWriter_fourcc(*'avc1')

    # 3. 预处理所有物体的帧范围 + 初始化写入器
    total_objects = len(obj_dict)
    print(f"总共需要处理 {total_objects} 个物体")

    writers = {}  # 存储所有视频写入器
    obj_ranges = {}  # 存储每个物体的起止帧

    for track_id, track_data in obj_dict.items():
        start_idx = max(0, track_data["start_frame"] - extra_frames)
        end_idx = min(total_frames - 1, track_data["end_frame"] + extra_frames)

        # 保证最小长度
        if end_idx - start_idx + 1 < min_frames:
            end_idx = start_idx + min_frames - 1

        output_path = os.path.join(output_dir, f"obj_{int(track_id):03d}.mp4")
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

        if not out.isOpened():
            raise Exception(f"无法初始化视频写入器: {output_path}")

        writers[track_id] = out
        obj_ranges[track_id] = (start_idx, end_idx)
        print(f"物体 {track_id}: 帧 {start_idx} ~ {end_idx} -> {output_path}")

    # 4. 单轮遍历视频，一次性写入所有需要的帧（最高效）
    frame_idx = 0
    print_interval = 50

    while cap.isOpened():
        # 读取一帧
        ret, frame = cap.read()
        if not ret:
            break

        # 遍历所有物体，判断当前帧是否需要写入
        for track_id, (start, end) in obj_ranges.items():
            if start <= frame_idx <= end:
                writers[track_id].write(frame)

        # 进度打印
        if frame_idx % print_interval == 0:
            progress = (frame_idx / total_frames) * 100
            print(f"处理帧: {frame_idx}/{total_frames} ({progress:.1f}%)")

        frame_idx += 1

    # 5. 释放所有写入器
    for track_id, out in writers.items():
        out.release()
        start, end = obj_ranges[track_id]
        print(f"物体 {track_id} 完成 | 总帧数: {end - start + 1}")

    # 6. 释放资源
    cap.release()
    cv2.destroyAllWindows()
    print(f"\n✅ 所有物体视频生成完成！目录: {output_dir}")

def process_track_id(
    track_id: int,
    frame_list: list[tuple[int, list[int]]],
    input_video_path: str,
    output_video_root: str,
    frame_width: int,
    frame_height: int,
    target_fps: int,
    frame_interval: int
) -> str:
    """
    处理单个track_id，生成对应的视频
    """
    # 计算需要生成的总帧数（确保覆盖所有物体帧且不少于两秒）
    min_frames_for_2s = 25 * 2  # 2秒 @ 25fps
    object_based_frames = len(frame_list) * frame_interval
    max_output_frame = max(object_based_frames, min_frames_for_2s)

    # 创建输出视频路径
    output_path = os.path.join(output_video_root, f"{track_id}.mp4")

    # 尝试使用GPU硬件编码
    try:
        # 对于不同平台的GPU编码，使用不同的fourcc
        # Windows平台使用h264_nvenc或h264_amf
        # 如果GPU编码不可用，会回退到CPU编码
        fourcc = cv2.VideoWriter_fourcc(*'h264') # type: ignore
        out = cv2.VideoWriter(output_path, fourcc, target_fps, (frame_width, frame_height))
        # 检查是否成功打开
        if not out.isOpened():
            # 尝试其他编码方式
            fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore
            out = cv2.VideoWriter(output_path, fourcc, target_fps, (frame_width, frame_height))
    except Exception:
        # 异常时使用CPU编码
        fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore
        out = cv2.VideoWriter(output_path, fourcc, target_fps, (frame_width, frame_height))

    if not out.isOpened():
        print(f"无法创建视频文件: {output_path}")
        return f"失败: {output_path}"

    # 打开原视频（每个线程独立打开，避免线程安全问题）
    cap = cv2.VideoCapture(input_video_path)
    if not cap.isOpened():
        out.release()
        return f"失败: 无法打开视频 {input_video_path}"

    # 生成视频帧
    current_output_frame = 0
    obj_frame_idx = 0

    while current_output_frame < max_output_frame:
        # 检查当前输出帧是否是5的倍数
        if current_output_frame % frame_interval == 0 and obj_frame_idx < len(frame_list):
            # 这是需要放置物体帧的位置
            original_frame_id, xyxy = frame_list[obj_frame_idx]

            # 设置原视频读取位置
            cap.set(cv2.CAP_PROP_POS_FRAMES, original_frame_id)
            ret, frame = cap.read()

            if not ret:
                # 读取失败，使用黑色帧
                output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
            else:
                # 有数据，截取对应区域
                x1, y1, x2, y2 = map(int, xyxy)
                # 确保坐标在有效范围内
                x1 = max(0, min(x1, frame_width))
                y1 = max(0, min(y1, frame_height))
                x2 = max(0, min(x2, frame_width))
                y2 = max(0, min(y2, frame_height))

                # 创建黑色背景
                output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
                # 将截取的区域放到输出帧中（保持原位置）
                if x2 > x1 and y2 > y1:
                    cropped = frame[y1:y2, x1:x2]
                    output_frame[y1:y2, x1:x2] = cropped

            # 移到下一个物体帧
            obj_frame_idx += 1
        else:
            # 剩余帧留黑
            output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)

        # 写入帧
        out.write(output_frame)
        current_output_frame += 1

    # 释放资源
    out.release()
    cap.release()
    print(f"已生成视频: {output_path}, 共 {current_output_frame} 帧")
    return f"成功: {output_path}"


def frame_all_to_obj_vid(
    json_data: dict,
    input_video_path: str,
    output_video_root: str,
) -> None:
    """
    根据标注数据从原视频中截取物体，生成ai读取专用视频

    参数:
        json_data: 标注数据
        input_video_path: 原视频路径
        output_video_root: 输出视频根目录
    """
    # 确保输出目录存在
    os.makedirs(output_video_root, exist_ok=True)

    # 1. 从 json_data 中提取数据，按 track_id 组织
    track_dict: dict[int, list[tuple[int, list[int]]]] = {}
    # 遍历每一帧
    for frame_id_str, detections in json_data.items():
        frame_id = int(frame_id_str)
        for det in detections:
            track_id = det.get("track_id", -1)
            xyxy = det.get("xyxy", [0, 0, 0, 0])
            if track_id not in track_dict:
                track_dict[track_id] = []
            track_dict[track_id].append((frame_id, xyxy))

    # 2. 获取视频信息（只需要获取一次）
    temp_cap = cv2.VideoCapture(input_video_path)
    if not temp_cap.isOpened():
        raise ValueError(f"无法打开视频: {input_video_path}")

    frame_width = int(temp_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(temp_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    temp_cap.release()

    # 目标fps为25，每5帧取一帧（0, 5, 10...）
    target_fps = 25
    frame_interval = 5  # 每隔5帧取一帧

    # 3. 使用多线程并行处理多个track_id
    # 根据CPU核心数设置线程池大小
    max_workers = min(os.cpu_count() or 4, len(track_dict))
    print(f"使用 {max_workers} 个线程并行处理")

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 提交所有任务
        future_to_track = {
            executor.submit(
                process_track_id,
                track_id,
                frame_list,
                input_video_path,
                output_video_root,
                frame_width,
                frame_height,
                target_fps,
                frame_interval
            ):
            track_id for track_id, frame_list in track_dict.items()
        }

        # 等待所有任务完成
        for future in as_completed(future_to_track):
            track_id = future_to_track[future]
            try:
                result = future.result()
                print(f"Track ID {track_id}: {result}")
            except Exception as e:
                print(f"Track ID {track_id} 处理失败: {str(e)}")

# def frame_all_to_obj_vid(
#     json_data: dict,
#     input_video_path: str,
#     output_video_root: str,
# ) -> None:
#     """
#     根据标注数据从原视频中截取物体，生成ai读取专用视频

#     参数:
#         json_data: 标注数据
#         input_video_path: 原视频路径
#         output_video_root: 输出视频根目录
#     """
#     # 确保输出目录存在
#     os.makedirs(output_video_root, exist_ok=True)

#     # 1. 从 json_data 中提取数据，按 track_id 组织
#     track_dict: dict[int, list[tuple[int, list[int]]]] = {}
#     # 遍历每一帧
#     for frame_id_str, detections in json_data.items():
#         frame_id = int(frame_id_str)
#         for det in detections:
#             track_id = det.get("track_id", -1)
#             xyxy = det.get("xyxy", [0, 0, 0, 0])
#             if track_id not in track_dict:
#                 track_dict[track_id] = []
#             track_dict[track_id].append((frame_id, xyxy))

#     # 2. 打开原视频
#     cap = cv2.VideoCapture(input_video_path)
#     if not cap.isOpened():
#         raise ValueError(f"无法打开视频: {input_video_path}")

#     # 获取原视频信息
#     original_fps = cap.get(cv2.CAP_PROP_FPS)
#     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
#     frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
#     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

#     # 目标fps为25，每5帧取一帧（0, 5, 10...）
#     target_fps = 25
#     frame_interval = 5  # 每隔5帧取一帧

#     # 为每个 track_id 生成视频
#     for track_id, frame_list in track_dict.items():
#         # 计算需要生成的总帧数（确保覆盖所有物体帧且不少于两秒）
#         min_frames_for_2s = 25 * 2  # 2秒 @ 25fps
#         object_based_frames = len(frame_list) * frame_interval
#         max_output_frame = max(object_based_frames, min_frames_for_2s)

#         # 创建输出视频路径
#         output_path = os.path.join(output_video_root, f"{track_id}.mp4")

#         # 创建视频写入器
#         fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore
#         out = cv2.VideoWriter(output_path, fourcc, target_fps, (frame_width, frame_height))

#         if not out.isOpened():
#             print(f"无法创建视频文件: {output_path}")
#             continue

#         # 生成视频帧
#         current_output_frame = 0
#         obj_frame_idx = 0

#         while current_output_frame < max_output_frame:
#             # 检查当前输出帧是否是5的倍数
#             if current_output_frame % frame_interval == 0 and obj_frame_idx < len(frame_list):
#                 # 这是需要放置物体帧的位置
#                 original_frame_id, xyxy = frame_list[obj_frame_idx]

#                 # 设置原视频读取位置
#                 cap.set(cv2.CAP_PROP_POS_FRAMES, original_frame_id)
#                 ret, frame = cap.read()

#                 if not ret:
#                     # 读取失败，使用黑色帧
#                     output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
#                 else:
#                     # 有数据，截取对应区域
#                     x1, y1, x2, y2 = map(int, xyxy)
#                     # 确保坐标在有效范围内
#                     x1 = max(0, min(x1, frame_width))
#                     y1 = max(0, min(y1, frame_height))
#                     x2 = max(0, min(x2, frame_width))
#                     y2 = max(0, min(y2, frame_height))

#                     # 创建黑色背景
#                     output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
#                     # 将截取的区域放到输出帧中（保持原位置）
#                     if x2 > x1 and y2 > y1:
#                         cropped = frame[y1:y2, x1:x2]
#                         output_frame[y1:y2, x1:x2] = cropped

#                 # 移到下一个物体帧
#                 obj_frame_idx += 1
#             else:
#                 # 剩余帧留黑
#                 output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)

#             # 写入帧
#             out.write(output_frame)
#             current_output_frame += 1

#         # 释放视频写入器
#         out.release()
#         print(f"已生成视频: {output_path}, 共 {current_output_frame} 帧")

#     # 释放原视频
#     cap.release()


def create_mian_vid_for_ai(
    input_video_path: str,
    output_folder: str
) -> str:
    """
    将原始视频的第0,1,2...帧映射到新视频的0,5,10...帧，其他帧留黑

    参数:
        input_video_path: 原始视频路径
        output_folder: 输出文件夹路径
    返回:
        str: 输出视频路径
    """
    # 确保输出目录存在
    os.makedirs(output_folder, exist_ok=True)

    # 构建输出视频路径
    output_video_path = os.path.join(output_folder, "mian_vid_ai.mp4")

    # 打开原视频
    cap = cv2.VideoCapture(input_video_path)
    if not cap.isOpened():
        raise ValueError(f"无法打开视频: {input_video_path}")

    # 获取原视频信息
    original_fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # 目标fps为25，每5帧取一帧（0, 5, 10...）
    target_fps = 25
    frame_interval = 5

    # 计算输出视频的总帧数
    # 确保覆盖所有原始帧且不少于两秒
    min_frames_for_2s = 25 * 2  # 2秒 @ 25fps
    max_output_frame = max(total_frames * frame_interval, min_frames_for_2s)

    # 创建视频写入器
    fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore
    out = cv2.VideoWriter(output_video_path, fourcc, target_fps, (frame_width, frame_height))

    if not out.isOpened():
        raise ValueError(f"无法创建视频文件: {output_video_path}")

    # 生成视频帧
    current_output_frame = 0
    original_frame_idx = 0

    while current_output_frame < max_output_frame:
        # 检查当前输出帧是否是5的倍数
        if current_output_frame % frame_interval == 0 and original_frame_idx < total_frames:
            # 这是需要放置原始帧的位置
            # 设置原视频读取位置
            cap.set(cv2.CAP_PROP_POS_FRAMES, original_frame_idx)
            ret, frame = cap.read()

            if not ret:
                # 读取失败，使用黑色帧
                output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
            else:
                # 有数据，使用原始帧
                output_frame = frame

            # 移到下一个原始帧
            original_frame_idx += 1
        else:
            # 剩余帧留黑
            output_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)

        # 写入帧
        out.write(output_frame)
        current_output_frame += 1

    # 释放资源
    cap.release()
    out.release()
    print(f"已生成视频: {output_video_path}, 共 {current_output_frame} 帧")

    return output_video_path