import base64
import tempfile
from pathlib import Path
from typing import Union
import cv2
import httpx


def _resize_frame(frame, max_short_edge: int):
    """按比例缩放帧，确保最短边不超过指定值"""
    h, w = frame.shape[:2]
    short_edge = min(h, w)
    if short_edge > max_short_edge:
        scale = max_short_edge / short_edge
        new_w = int(w * scale)
        new_h = int(h * scale)
        frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
    return frame


def contents(video_source: Union[str, Path], prompt_start: str = None, prompt_end=None, video_name: str = None,
             fps: float = 1.0,
             max_frames: int = 10,
             max_short_edge: int = 768,
             sampling_mode: str = "uniform") -> list[str]:
    """
    从视频中提取帧并构建 LLM 消息内容。

    Args:
        video_source: 视频源，可以是文件路径或 URL。
        prompt_start: 提示词的开始部分。
        prompt_end: 提示词的结束部分。
        video_name: 视频名称。
        fps: 提取帧的帧率。
        max_frames: 最大帧数。
        max_short_edge: 最大短边长度。
        sampling_mode: 当帧数超过 max_frames 时的采样策略。
            - "uniform": 均匀采样（默认）
            - "head": 保留前面的帧，抛弃后面的帧
    """
    source_str = str(video_source)
    temp_file = None

    try:
        # 如果是 URL，先下载到临时文件
        if source_str.startswith(("http://", "https://")):
            resp = httpx.get(source_str, timeout=60.0)
            resp.raise_for_status()
            temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
            temp_file.write(resp.content)
            temp_file.close()
            video_path = temp_file.name
        else:
            video_path = str(video_source)
            if not Path(video_path).exists():
                raise FileNotFoundError(f"视频文件不存在: {video_path}")

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

        # 获取视频原始帧率
        video_fps = cap.get(cv2.CAP_PROP_FPS)

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

        frames_base64 = []
        frame_count = 0

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

            # 按指定间隔提取帧
            if frame_count % frame_interval == 0:
                # 缩放帧以控制最短边长度
                frame = _resize_frame(frame, max_short_edge)
                # 编码为 JPEG 并转 base64
                _, buffer = cv2.imencode('.jpg', frame)
                frame_base64 = base64.b64encode(buffer).decode('utf-8')
                frames_base64.append(frame_base64)

            frame_count += 1

        cap.release()

        if len(frames_base64) > max_frames:
            import math
            step = len(frames_base64) / max_frames
            sampled = [frames_base64[min(int(i * step), len(frames_base64) - 1)] for i in range(max_frames)]
            frames_base64 = sampled

        # 构建消息内容：提示词 + 所有帧图片
        video_prompt = (
            f"以下是从视频({video_name})中按时间顺序提取的 {len(frames_base64)} 帧画面，视频原始帧率为 {video_fps:.2f} fps，"
            f"抽帧间隔为 {frame_interval} 帧（约每 {frame_interval / video_fps:.2f} 秒一帧），请将它们视为一个连续的视频进行分析。"
        )
        content = []
        if prompt_start is not None:
            content.append({"type": "text", "text": prompt_start})
        content.append({"type": "text", "text": video_prompt})

        for frame_base64 in frames_base64:
            content.append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/jpeg;base64,{frame_base64}"
                }
            })
        if prompt_end is not None:
            content.append({"type": "text", "text": prompt_end})
        return content

    finally:
        # 清理临时文件
        if temp_file and Path(temp_file.name).exists():
            Path(temp_file.name).unlink()