#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
视频上传Function

功能: 下载视频到本地并上传至Gemini，保存上传链接到state
"""

import os
import time
import tempfile
import requests
from pathlib import Path
from urllib.parse import urlparse
from typing import Dict, Any, Optional, Tuple

from google import genai

from src.components.functions.base import BaseFunction
from src.states.what_deconstruction_state import WhatDeconstructionState
from src.utils.logger import get_logger
from src.utils.llm_invoker import LLMInvoker

logger = get_logger(__name__)


class VideoUploadFunction(BaseFunction[Dict[str, Any], Dict[str, Any]]):
    """视频上传函数
    
    功能：
    - 从URL下载视频到本地
    - 上传视频到Gemini File API
    - 保存上传后的文件URI到state
    """

    def __init__(
        self,
        name: str = "video_upload_function",
        description: str = "下载视频并上传至Gemini，保存上传链接"
    ):
        super().__init__(name, description)

    def execute(
        self,
        input_data: Dict[str, Any],
        context: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """执行视频上传
        
        Args:
            input_data: 包含video字段的状态字典
            context: 上下文信息
            
        Returns:
            更新后的状态字典，包含video_uploaded_uri字段
        """
        try:
            video_url = input_data.get("video", "")
            
            if not video_url:
                logger.warning("未提供视频URL，跳过上传")
                return {
                    **input_data,
                    "video_uploaded_uri": None,
                    "video_upload_error": "未提供视频URL"
                }

            logger.info(f"开始下载视频: {video_url}")

            # 1. 下载视频到本地临时文件（或使用examples目录下的现有文件）
            # 从input_data中获取channel_content_id，用于查找examples目录下的文件
            channel_content_id = input_data.get("channel_content_id", "")
            local_video_path, is_temp_file = self._download_video(video_url, channel_content_id)
            
            if not local_video_path:
                return {
                    **input_data,
                    "video_uploaded_uri": None,
                    "video_upload_error": "视频下载失败"
                }

            logger.info(f"视频文件路径: {local_video_path}")

            # 2. 上传视频到Gemini（使用新的API，带state校验）
            video_file = LLMInvoker.upload_video_to_gemini(local_video_path)
            
            # # 3. 清理临时文件（只有临时文件才需要清理）
            # if is_temp_file:
            #     try:
            #         os.remove(local_video_path)
            #         logger.info(f"临时文件已删除: {local_video_path}")
            #     except Exception as e:
            #         logger.warning(f"删除临时文件失败: {e}")
            # else:
            #     logger.info(f"使用examples目录下的文件，不删除: {local_video_path}")

            if not video_file:
                return {
                    **input_data,
                    "video_uploaded_uri": None,
                    "video_file_name": None,
                    "video_upload_error": "视频上传到Gemini失败"
                }

            # 获取文件URI和名称
            file_uri = None
            file_name = None
            if hasattr(video_file, 'uri'):
                file_uri = video_file.uri
            elif hasattr(video_file, 'name'):
                file_name = video_file.name
                file_uri = f"https://generativelanguage.googleapis.com/v1beta/files/{file_name}"
            
            logger.info(f"视频上传成功，文件名称: {file_name}")

            # 4. 更新state
            return {
                **input_data,
                "video_uploaded_uri": file_uri,  # 兼容旧版本
                "video_file_name": file_name,  # 新字段，用于获取文件对象
                "video_upload_error": None
            }

        except Exception as e:
            logger.error(f"视频上传失败: {e}", exc_info=True)
            return {
                **input_data,
                "video_uploaded_uri": None,
                "video_file_name": None,
                "video_upload_error": str(e)
            }

    def _download_video(self, video_url: str, channel_content_id: str = "") -> Tuple[Optional[str], bool]:
        """下载视频到本地临时文件，或使用examples目录下的现有文件
        
        Args:
            video_url: 视频URL
            channel_content_id: 频道内容ID，用于查找examples目录下的文件
            
        Returns:
            (本地文件路径, 是否为临时文件) 的元组，失败返回 (None, True)
            如果使用examples目录下的文件，返回 (文件路径, False)
            如果下载到临时文件，返回 (临时文件路径, True)
        """
        try:
            # 1. 首先检查examples目录下是否有对应的mp4文件
            existing_file = self._check_examples_directory(channel_content_id)
            if existing_file:
                logger.info(f"在examples目录下找到现有文件，直接使用: {existing_file}")
                return existing_file, False
            
            # 2. 如果没有找到，则下载到临时文件
            logger.info("未在examples目录下找到同名文件，开始下载...")
            
            # 创建临时文件
            temp_dir = tempfile.gettempdir()
            temp_file = tempfile.NamedTemporaryFile(
                dir=temp_dir,
                suffix=".mp4",
                delete=False
            )
            temp_path = temp_file.name
            temp_file.close()

            # 下载视频（带重试机制）
            max_retries = 3
            retry_count = 0
            last_exception = None
            
            while retry_count < max_retries:
                try:
                    if retry_count > 0:
                        logger.info(f"重试下载视频 (第 {retry_count}/{max_retries-1} 次)...")
                    
                    # 使用 Session 进行下载（参考提供的代码）
                    session = requests.Session()
                    session.headers.update({
                        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"# 禁用压缩，避免IncompleteRead问题
                    })
                    
                    # 下载视频（增加超时时间）
                    response = session.get(
                        video_url,
                        timeout=120,  # (连接超时, 读取超时)
                        stream=True
                    )
                    response.raise_for_status()

                    # 确保目录存在
                    os.makedirs(os.path.dirname(temp_path), exist_ok=True)

                    # 写入文件
                    with open(temp_path, "wb") as f:
                        for chunk in response.iter_content(chunk_size=8192):
                            if chunk:
                                f.write(chunk)
                    
                    # 验证文件大小
                    file_size = os.path.getsize(temp_path)
                    if file_size == 0:
                        raise ValueError("下载的文件大小为0")
                    
                    logger.info(f"视频下载完成，大小: {file_size / 1024 / 1024:.2f} MB")
                    return temp_path, True
                    
                except (requests.exceptions.ChunkedEncodingError, 
                        requests.exceptions.ConnectionError,
                        requests.exceptions.Timeout,
                        requests.exceptions.RequestException) as e:
                    last_exception = e
                    retry_count += 1
                    
                    # 清理不完整的文件
                    if os.path.exists(temp_path):
                        try:
                            os.remove(temp_path)
                        except:
                            pass
                    
                    if retry_count < max_retries:
                        wait_time = retry_count * 2  # 递增等待时间：2秒、4秒
                        logger.warning(f"下载失败 (尝试 {retry_count}/{max_retries}): {e}")
                        logger.info(f"等待 {wait_time} 秒后重试...")
                        time.sleep(wait_time)
                    else:
                        logger.error(f"下载失败，已重试 {max_retries} 次: {e}")
                        raise
                except Exception as e:
                    # 其他类型的异常直接抛出，不重试
                    if os.path.exists(temp_path):
                        try:
                            os.remove(temp_path)
                        except:
                            pass
                    raise
            
            # 如果所有重试都失败了
            if last_exception:
                raise last_exception

        except Exception as e:
            logger.error(f"下载视频失败: {e}", exc_info=True)
            return None, True
    
    def _check_examples_directory(self, channel_content_id: str) -> Optional[str]:
        """检查examples目录下是否有对应的mp4文件
        
        文件路径格式：examples/{channel_content_id}/{channel_content_id}.mp4
        
        Args:
            channel_content_id: 频道内容ID（目录名）
            
        Returns:
            如果找到文件，返回文件路径；否则返回None
        """
        try:
            # 如果没有提供channel_content_id，无法查找
            if not channel_content_id:
                logger.info("未提供channel_content_id，跳过examples目录检查")
                return None
            
            # 获取项目根目录
            # __file__ 是 src/components/functions/video_upload_function.py
            # 需要往上4层才能到项目根目录
            project_root = Path(__file__).parent.parent.parent.parent
            examples_dir = project_root / "examples"
            
            if not examples_dir.exists():
                logger.info(f"examples目录不存在: {examples_dir}")
                return None
            
            # 构建文件路径：examples/{channel_content_id}/{channel_content_id}.mp4
            video_dir = examples_dir / channel_content_id
            mp4_file = video_dir / f"{channel_content_id}.mp4"
            logger.info(f"构建文件路径: {mp4_file}")
            # 检查文件是否存在
            if mp4_file.exists() and mp4_file.is_file():
                logger.info(f"在examples目录下找到文件: {mp4_file}")
                return str(mp4_file)
            
            logger.debug(f"在examples目录下未找到文件: {mp4_file}")
            return None
            
        except Exception as e:
            logger.warning(f"检查examples目录时出错: {e}", exc_info=True)
            return None


