import base64
import hashlib
import subprocess
from pathlib import Path

import cv2

try:
    from .m3u8_to_mp4_sei import (
        download_m3u8_to_mp4,
        extract_h264_es_from_mp4,
        parse_h264_sei,
        extract_utc_from_sei,
    )
except ImportError:
    from m3u8_to_mp4_sei import (
        download_m3u8_to_mp4,
        extract_h264_es_from_mp4,
        parse_h264_sei,
        extract_utc_from_sei,
    )

# 用于区分传入的 start_utc/end_utc 是绝对 UTC 时间戳还是秒偏移
_IS_ABSOLUTE_UTC_THRESHOLD = 1_000_000_000


class Video2Frame:
    def __init__(self, cache_dir):
        self.cache_dir = Path(cache_dir)
        self.video_dir = self.cache_dir / "videos"
        self.frames_dir = self.cache_dir / "frames"
        self.video_dir.mkdir(parents=True, exist_ok=True)
        self.frames_dir.mkdir(parents=True, exist_ok=True)

    def to_frames(self, url, start_utc, end_utc, fps, roi=None, max_px_area=None) -> list:
        """
        下载 m3u8 视频流的指定片段并提取帧，保存到 cache_dir 下。

        - video_dir 中保存的是 start_utc ~ end_utc 截取后的视频片段（而非完整视频）。
        - 如果视频中包含 SEI UTC 信息且传入的 start_utc/end_utc 为绝对时间戳，
          则会基于 SEI UTC 进行帧级定位；否则将 start_utc/end_utc 视为秒偏移量。

        Args:
            url: m3u8 视频流地址。
            start_utc: 截取开始时间（秒偏移或绝对 UTC 时间戳）。
            end_utc: 截取结束时间（秒偏移或绝对 UTC 时间戳）。
            fps: 目标抽帧帧率。
            roi: 感兴趣区域 (x, y, w, h)，可选。
            max_px_area: 最大像素面积，超过则等比例缩小，可选。

        Returns:
            提取的帧图片路径列表。
        """
        if end_utc <= start_utc:
            raise ValueError("end_utc 必须大于 start_utc")

        unique_str = f"{url}_{start_utc}_{end_utc}"
        unique_id = hashlib.md5(unique_str.encode("utf-8")).hexdigest()

        clip_path = self.video_dir / f"{unique_id}.mp4"
        frame_output_dir = self.frames_dir / unique_id
        frame_output_dir.mkdir(parents=True, exist_ok=True)

        # 判断是否为绝对 UTC 时间戳
        is_absolute_utc = (
            start_utc > _IS_ABSOLUTE_UTC_THRESHOLD
            and end_utc > _IS_ABSOLUTE_UTC_THRESHOLD
        )

        # 若片段未缓存，先截取目标片段再保存
        if not clip_path.exists():
            if is_absolute_utc:
                # 绝对 UTC 模式：先下载完整视频到临时文件，解析 SEI 后截取片段
                temp_path = self.video_dir / f"{unique_id}_full.mp4"
                download_m3u8_to_mp4(url, str(temp_path))

                # 提取 SEI UTC 信息
                utc_records = []
                try:
                    es_data = extract_h264_es_from_mp4(str(temp_path))
                    sei_list = parse_h264_sei(es_data)
                    utc_records = extract_utc_from_sei(sei_list)
                except Exception:
                    pass

                if utc_records:
                    utc_list = [r["utc"] for r in utc_records]
                    start_frame_idx = self._find_frame_idx(
                        utc_list, start_utc, find_last_not_greater=False
                    )
                    end_frame_idx = self._find_frame_idx(
                        utc_list, end_utc, find_last_not_greater=True
                    )

                    cap_temp = cv2.VideoCapture(str(temp_path))
                    video_fps = cap_temp.get(cv2.CAP_PROP_FPS)
                    cap_temp.release()
                    if video_fps <= 0:
                        video_fps = 25.0

                    start_sec = start_frame_idx / video_fps
                    duration = (end_frame_idx - start_frame_idx) / video_fps
                    self._ffmpeg_extract_clip(
                        str(temp_path), str(clip_path), start_sec, duration
                    )
                else:
                    # 无 SEI 无法定位，直接将完整视频重命名为片段
                    temp_path.rename(clip_path)

                # 清理临时完整视频
                if temp_path.exists():
                    temp_path.unlink(missing_ok=True)
            else:
                # 秒偏移模式：直接用 ffmpeg 从 URL 截取片段
                duration = end_utc - start_utc
                self._ffmpeg_download_clip(
                    url, str(clip_path), start_utc, duration
                )

        # 从截取后的片段中提取帧
        cap = cv2.VideoCapture(str(clip_path))
        if not cap.isOpened():
            raise ValueError(f"无法打开视频片段: {clip_path}")

        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        video_fps = cap.get(cv2.CAP_PROP_FPS)
        if video_fps <= 0:
            video_fps = 25.0

        # 计算抽帧间隔
        frame_interval = int(video_fps / fps) if fps > 0 else int(video_fps)
        if frame_interval < 1:
            frame_interval = 1

        frame_paths = []
        saved_count = 0
        frame_count = 0

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            if frame_count % frame_interval == 0:
                if roi is not None:
                    x, y, w, h = roi
                    frame = frame[y:y + h, x:x + w]

                if max_px_area is not None:
                    h, w = frame.shape[:2]
                    area = h * w
                    if area > max_px_area:
                        scale = (max_px_area / area) ** 0.5
                        new_w = int(w * scale)
                        new_h = int(h * scale)
                        frame = cv2.resize(
                            frame, (new_w, new_h), interpolation=cv2.INTER_AREA
                        )

                frame_path = frame_output_dir / f"frame_{saved_count:06d}.jpg"
                cv2.imwrite(str(frame_path), frame)
                frame_paths.append(str(frame_path))
                saved_count += 1

            frame_count += 1

        cap.release()
        return frame_paths

    @staticmethod
    def _find_frame_idx(utc_list, target_utc, find_last_not_greater=True):
        """
        在 utc_list 中查找最接近 target_utc 的帧索引。

        Args:
            utc_list: 按帧顺序排列的 UTC 列表。
            target_utc: 目标 UTC。
            find_last_not_greater: True 返回最后一个 <= target_utc 的索引；
                                   False 返回第一个 >= target_utc 的索引。
        """
        if not utc_list:
            return 0

        if find_last_not_greater:
            idx = 0
            for i, utc in enumerate(utc_list):
                if utc > target_utc:
                    break
                idx = i
            return idx
        else:
            for i, utc in enumerate(utc_list):
                if utc >= target_utc:
                    return i
            return len(utc_list) - 1

    @staticmethod
    def _ffmpeg_extract_clip(input_path, output_path, start_sec, duration):
        """使用 ffmpeg 从本地视频截取指定片段。"""
        cmd = [
            "ffmpeg", "-y",
            "-i", input_path,
            "-ss", str(start_sec),
            "-t", str(duration),
            "-c", "copy",
            "-bsf:a", "aac_adtstoasc",
            "-movflags", "+faststart",
            output_path,
        ]
        try:
            subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        except subprocess.CalledProcessError as e:
            raise RuntimeError(
                f"ffmpeg 截取片段失败: {e.stderr.decode('utf-8', errors='ignore')}"
            )

    @staticmethod
    def _ffmpeg_download_clip(url, output_path, start_sec, duration):
        """使用 ffmpeg 从 URL 下载并截取指定片段。"""
        cmd = [
            "ffmpeg", "-y",
            "-ss", str(start_sec),
            "-fflags", "+discardcorrupt",
            "-i", url,
            "-t", str(duration),
            "-c", "copy",
            "-bsf:a", "aac_adtstoasc",
            "-movflags", "+faststart",
            output_path,
        ]
        try:
            subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        except subprocess.CalledProcessError as e:
            raise RuntimeError(
                f"ffmpeg 下载片段失败: {e.stderr.decode('utf-8', errors='ignore')}"
            )


def frames2content(frames):
    """
    将帧列表（图片路径列表）转为 LLM content 列表。
    格式与 llm_video_content.py 中的 contents 返回的 content 一致。
    """
    content = []
    video_prompt = (
        f"以下是从视频中按时间顺序提取的 {len(frames)} 帧画面，"
        f"请将它们视为一个连续的视频进行分析。"
    )
    content.append({"type": "text", "text": video_prompt})

    for frame_path in frames:
        with open(frame_path, "rb") as f:
            img_bytes = f.read()
        b64_str = base64.b64encode(img_bytes).decode("utf-8")
        content.append({
            "type": "image_url",
            "image_url": {
                "url": f"data:image/jpeg;base64,{b64_str}"
            }
        })

    return content


if __name__ == "__main__":
    import argparse
    import sys
    import tempfile

    parser = argparse.ArgumentParser(description="Video2Frame 测试脚本")
    parser.add_argument(
        "--url",
        default=r"D:\pythonProject\learn\3b3e99cf4ca84c3782503d8817242de2.mp4",
        help="测试用 m3u8 地址（默认使用公开测试流）",
    )
    parser.add_argument(
        "--cache-dir",
        default=None,
        help="缓存目录（默认自动创建临时目录）",
    )
    parser.add_argument(
        "--start",
        type=float,
        default=0,
        help="截取开始时间，单位：秒（默认 0）",
    )
    parser.add_argument(
        "--end",
        type=float,
        default=10,
        help="截取结束时间，单位：秒（默认 10）",
    )
    parser.add_argument(
        "--fps",
        type=float,
        default=1.0,
        help="抽帧帧率（默认 1.0）",
    )
    args = parser.parse_args()

    cache_dir = args.cache_dir or tempfile.mkdtemp(prefix="video2frame_test_")
    print(f"[INFO] 缓存目录: {cache_dir}")

    v2f = Video2Frame(cache_dir)

    # TEST 1: 基本抽帧
    print(
        f"\n[TEST 1] to_frames: url={args.url}, "
        f"start={args.start}s, end={args.end}s, fps={args.fps}"
    )
    try:
        frames = v2f.to_frames(
            url=args.url,
            start_utc=args.start,
            end_utc=args.end,
            fps=args.fps,
        )
        print(f"  成功提取 {len(frames)} 帧")
        if frames:
            print(f"  首帧: {frames[0]}")
            print(f"  末帧: {frames[-1]}")
    except Exception as e:
        print(f"  失败: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)

    # TEST 2: frames2content
    print("\n[TEST 2] frames2content")
    content = frames2content(frames)
    print(f"  content 列表长度: {len(content)}")
    for item in content[:3]:
        preview = (
            item["text"][:60] + "..."
            if item["type"] == "text"
            else item["image_url"]["url"][:60] + "..."
        )
        print(f"  - type={item['type']}: {preview}")
    if len(content) > 3:
        print(f"  ... 还有 {len(content) - 3} 个元素")

    # TEST 3: ROI + max_px_area
    print("\n[TEST 3] to_frames with roi + max_px_area")
    try:
        end_crop = min(args.start + 5, args.end)
        frames_cropped = v2f.to_frames(
            url=args.url,
            start_utc=args.start,
            end_utc=end_crop,
            fps=1.0,
            roi=None,
            max_px_area=200_000,
        )
        print(f"  成功提取 {len(frames_cropped)} 帧（带 ROI 裁剪和面积缩放）")
    except Exception as e:
        print(f"  跳过/失败: {e}")

    # TEST 4: SEI UTC 模式说明（仅演示，需真实含 SEI 的流才能运行）
    print("\n[TEST 4] SEI UTC 绝对时间戳模式（示例代码，需替换为含 SEI 的流）")
    print(
        "  # 示例：假设视频第一帧 UTC 为 1715600000，提取 10 秒片段\n"
        "  # frames = v2f.to_frames(url, start_utc=1715600000, end_utc=1715600010, fps=1)"
    )

    print("\n[INFO] 所有测试完成")