#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
LLM调用封装工具类

功能: 统一封装LLM调用逻辑,提供安全的调用接口
"""

import json
import os
import time
from typing import Any, List, Dict, Optional

from google import genai

from src.utils.logger import get_logger
from src.utils.json_extractor import JSONExtractor

logger = get_logger(__name__)

# Gemini 文件处理相关常量
FILE_PROCESS_TIMEOUT = 300  # 文件处理超时时间（秒）：5分钟
API_RETRY_INTERVAL = 5  # API重试间隔（秒）

# 全局 client 实例
_client = None


def _get_client():
    """获取 Gemini Client 实例"""
    global _client
    if _client is None:
        api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
        if not api_key:
            raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY not found in environment")
        _client = genai.Client(api_key=api_key)
    return _client


def get_video_file_from_state(state: Dict[str, Any]) -> Optional["genai.types.File"]:
    """从 state 中获取视频文件对象
    
    Args:
        state: 状态字典，可能包含：
            - video_file: 直接存储的文件对象（优先）
            - video_file_name: 文件名称
            - video_uploaded_uri: 文件URI
        
    Returns:
        Gemini文件对象，失败返回None
    """
    try:
        # 优先检查是否直接存储了文件对象
        video_file = state.get("video_file")
        if video_file:
            logger.info("从state中直接获取视频文件对象")
            return video_file
        
        # 如果没有直接存储，尝试通过文件名或URI获取
        client = _get_client()
        
        # 优先使用 video_file_name
        file_name = state.get("video_file_name")
        if file_name:
            try:
                logger.info(f"通过 file_name 获取文件: {file_name}")
                return client.files.get(name=file_name)
            except Exception as e:
                logger.warning(f"通过 file_name 获取文件失败: {e}")
        
        # 如果没有 file_name，尝试从 URI 中提取
        video_uri = state.get("video_uploaded_uri", "")
        if video_uri:
            # 从 URI 中提取文件名称
            # URI 格式: https://generativelanguage.googleapis.com/v1beta/files/{file_name}
            if "/files/" in video_uri:
                file_name = video_uri.split("/files/")[-1]
                try:
                    logger.info(f"通过 URI 提取的 file_name 获取文件: {file_name}")
                    return client.files.get(name=file_name)
                except Exception as e:
                    logger.warning(f"通过 URI 提取的 file_name 获取文件失败: {e}")
        
        logger.error("无法从 state 中获取视频文件对象")
        return None
    except Exception as e:
        logger.error(f"获取视频文件对象失败: {e}", exc_info=True)
        return None


class LLMInvoker:
    """LLM调用封装类"""

    @staticmethod
    def wait_for_file_processing(uploaded_file: "genai.types.File") -> Optional["genai.types.File"]:
        """等待 Gemini 完成文件处理
        
        Args:
            uploaded_file: 已上传的文件对象
            
        Returns:
            处理完成后的文件对象，失败或超时返回None
        """
        start = time.time()
        current = uploaded_file
        client = _get_client()
        
        if not hasattr(current, 'state'):
            return current
        
        state_name = getattr(current.state, 'name', None) if hasattr(current.state, 'name') else str(current.state)
        
        if state_name != "PROCESSING":
            if state_name == "FAILED":
                logger.error(f"文件处理失败: {current.state}")
                return None
            elif state_name in ["ACTIVE", "COMPLETE", "READY"]:
                return current
        
        logger.info("等待文件处理...")
        
        while True:
            elapsed = time.time() - start
            if elapsed > FILE_PROCESS_TIMEOUT:
                logger.error(f"文件处理超时（超过 {FILE_PROCESS_TIMEOUT} 秒）")
                return None
            
            time.sleep(API_RETRY_INTERVAL)
            
            if hasattr(current, 'name'):
                try:
                    current = client.files.get(name=current.name)
                    if not hasattr(current, 'state'):
                        break
                    
                    state_name = getattr(current.state, 'name', None) if hasattr(current.state, 'name') else str(current.state)
                    
                    if state_name == "FAILED":
                        logger.error(f"文件处理失败: {current.state}")
                        return None
                    elif state_name in ["ACTIVE", "COMPLETE", "READY"]:
                        break
                except Exception as e:
                    logger.error(f"获取文件状态失败: {e}")
                    return None
            else:
                return None
        
        logger.info("文件处理完成")
        return current

    @staticmethod
    def upload_video_to_gemini(local_file_path: str) -> Optional["genai.types.File"]:
        """上传视频到Gemini并等待处理完成
        
        Args:
            local_file_path: 本地视频文件路径
            
        Returns:
            处理完成后的文件对象，失败返回None
        """
        try:
            client = _get_client()
            logger.info(f"上传视频到Gemini: {local_file_path}")
            
            uploaded_file = client.files.upload(file=local_file_path)
            
            processed_file = LLMInvoker.wait_for_file_processing(uploaded_file)
            if not processed_file:
                raise Exception("文件处理失败")
            
            logger.info("视频上传并处理完成")
            return processed_file
        except Exception as e:
            logger.error(f"上传视频到Gemini失败: {e}", exc_info=True)
            return None

    @staticmethod
    def safe_invoke(
        agent: Any,
        operation_name: str,
        messages: List[Dict[str, Any]],
        fallback: Any = None
    ) -> Any:
        """安全的LLM调用

        Args:
            agent: Agent实例,必须有model属性和is_initialized属性
            operation_name: 操作名称,用于日志记录
            messages: LLM消息列表
            fallback: 调用失败时返回的降级值

        Returns:
            解析后的JSON对象,失败时返回fallback
        """
        try:
            if not agent.is_initialized:
                agent.initialize()

            response = agent.model.invoke(messages)
            content = response.content.strip()

            result = JSONExtractor.extract_and_parse(content, {})

            if not result:
                logger.warning(
                    f"{operation_name} - LLM返回内容为空或解析失败\n"
                    f"原始内容: {content}..."
                )
                return fallback if fallback is not None else {}

            return result

        except json.JSONDecodeError as e:
            logger.error(
                f"{operation_name} - JSON解析失败: {e}\n"
                f"原始内容: {content if 'content' in locals() else 'N/A'}..."
            )
            return fallback if fallback is not None else {}

        except Exception as e:
            logger.error(f"{operation_name} - 调用失败: {e}", exc_info=True)
            return fallback if fallback is not None else {}

    @staticmethod
    def safe_invoke_video_analysis(
        operation_name: str,
        video_file: "genai.types.File",
        prompt: str,
        agent: Any = None,
        fallback: Any = None
    ) -> Any:
        """安全的视频分析LLM调用（使用新的Gemini API）
        
        Args:
            operation_name: 操作名称，用于日志记录
            video_file: Gemini文件对象（已上传并处理完成的文件）
            prompt: 分析提示词
            agent: Agent实例（可选，用于获取模型配置）
            fallback: 调用失败时返回的降级值
            
        Returns:
            解析后的JSON对象，失败时返回fallback
        """
        try:
            client = _get_client()
            
            # 获取模型名称（从agent获取，或使用默认值）
            model_name = "gemini-2.5-flash"
            if agent is not None:
                # 尝试从agent获取模型配置
                if hasattr(agent, 'model') and hasattr(agent.model, 'model_name'):
                    model_name = agent.model.model_name
                elif hasattr(agent, 'model_provider'):
                    # 如果agent有model_provider，使用默认模型
                    pass
            
            logger.info(f"{operation_name} - 开始视频分析...")
            
            # 调用模型进行视频分析（参考 test.py 的方式）
            response = client.models.generate_content(
                model=model_name,
                contents=[video_file, prompt]
            )
            
            # 提取响应内容
            if not response:
                logger.warning(f"{operation_name} - 响应为空")
                return fallback if fallback is not None else {}
            
            # 获取文本内容
            if hasattr(response, 'text'):
                content = response.text.strip()
            elif hasattr(response, 'content'):
                if isinstance(response.content, str):
                    content = response.content.strip()
                elif hasattr(response.content, 'text'):
                    content = response.content.text.strip()
                else:
                    content = str(response.content).strip()
            else:
                content = str(response).strip()
            
            if not content:
                logger.warning(f"{operation_name} - 响应内容为空")
                return fallback if fallback is not None else {}
            
            # 解析JSON
            result = JSONExtractor.extract_and_parse(content, {})
            
            if not result:
                logger.warning(
                    f"{operation_name} - LLM返回内容为空或解析失败\n"
                    f"原始内容: {content[:200]}..."
                )
                return fallback if fallback is not None else {}
            
            logger.info(f"{operation_name} - 视频分析完成")
            return result
            
        except json.JSONDecodeError as e:
            logger.error(
                f"{operation_name} - JSON解析失败: {e}\n"
                f"原始内容: {content if 'content' in locals() else 'N/A'}..."
            )
            return fallback if fallback is not None else {}
            
        except Exception as e:
            logger.error(f"{operation_name} - 视频分析调用失败: {e}", exc_info=True)
            return fallback if fallback is not None else {}
