"""
DashScope API 客户端

用于调用阿里云通义千问3-VL-Plus模型的API客户端。
使用 DashScope 原生 SDK 而不是 OpenAI 兼容 API。
"""

import os
from typing import Dict, Any, Optional
from dotenv import load_dotenv
import dashscope
from dashscope import MultiModalConversation

# 加载环境变量
load_dotenv()

# 图片生成的预设尺寸（qwen-image-plus 模型支持的尺寸）
IMAGE_SIZES = {
    '1664*928': '1664*928',      # 宽屏
    '1472*1140': '1472*1140',    # 宽屏
    '1328*1328': '1328*1328',    # 正方形（默认）
    '1140*1472': '1140*1472',    # 竖屏
    '928*1664': '928*1664',      # 竖屏
}


class DashScopeClient:
    """DashScope API 客户端"""

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: str = "qwen3-vl-plus"
    ):
        """
        初始化 DashScope 客户端

        Args:
            api_key: API 密钥，如果为 None 则从环境变量获取
            model: 模型名称
        """
        self.api_key = api_key or os.getenv("IMAGE_MCP_DASHSCOPE_API_KEY")
        if not self.api_key:
            raise ValueError(
                "DashScope API 密钥未设置。请设置环境变量 IMAGE_MCP_DASHSCOPE_API_KEY "
                "或在创建客户端时传入 api_key 参数"
            )

        # 设置全局 API 密钥
        dashscope.api_key = self.api_key
        self.model = model

    async def analyze_image(
        self,
        image_base64: str,
        mime_type: str,
        question: str = "请分析这张图片的内容",
        enable_thinking: bool = True,
        thinking_budget: int = 81920
    ) -> Dict[str, Any]:
        """
        分析图片内容

        Args:
            image_base64: base64编码的图片数据
            mime_type: 图片的MIME类型
            question: 关于图片的问题
            enable_thinking: 是否启用思考功能
            thinking_budget: 思考过程最大 Token 数

        Returns:
            Dict[str, Any]: API响应结果，包含 analysis、reasoning_content 等字段

        Raises:
            Exception: API调用失败
        """
        try:
            # 构建消息内容 (DashScope 原生格式)
            # 使用 data URI 格式传递 base64 图片
            image_url = f"data:{mime_type};base64,{image_base64}"

            messages = [
                {
                    "role": "user",
                    "content": [
                        {
                            "image": image_url
                        },
                        {
                            "text": question
                        }
                    ]
                }
            ]

            # 调用 MultiModalConversation API
            response = MultiModalConversation.call(
                model=self.model,
                messages=messages,
                stream=False,  # 非流式响应
                enable_thinking=enable_thinking,
                thinking_budget=thinking_budget
            )

            return self._parse_response(response)

        except Exception as e:
            raise Exception(f"图片分析失败: {str(e)}")

    def _parse_response(self, response: Any) -> Dict[str, Any]:
        """
        解析 API 响应

        Args:
            response: DashScope API 响应对象

        Returns:
            Dict[str, Any]: 解析后的响应数据
        """
        try:
            # 检查是否有错误
            if not response:
                return {
                    "success": False,
                    "error": "API 响应为空"
                }

            # 检查响应状态
            if hasattr(response, 'status_code') and response.status_code != 200:
                error_msg = getattr(response, 'message', f"HTTP {response.status_code}")
                return {
                    "success": False,
                    "error": f"API 错误: {error_msg}"
                }

            # 提取内容
            if hasattr(response, 'output') and response.output:
                output = response.output
                if hasattr(output, 'choices') and output.choices:
                    choice = output.choices[0]
                    if hasattr(choice, 'message'):
                        message = choice.message

                        # 提取思考过程（如果有）
                        reasoning_content = getattr(message, 'reasoning_content', '')

                        # 提取回复内容
                        content = getattr(message, 'content', [])
                        text_content = ""
                        if isinstance(content, list) and len(content) > 0:
                            # 提取第一个文本内容
                            for item in content:
                                if isinstance(item, dict) and 'text' in item:
                                    text_content = item['text']
                                    break
                                elif hasattr(item, 'text'):
                                    text_content = item.text
                                    break
                        elif isinstance(content, str):
                            text_content = content

                        # 构建返回结果
                        result = {
                            "success": True,
                            "analysis": text_content,
                            "model": self.model,
                        }

                        # 如果有思考过程，添加到结果
                        if reasoning_content:
                            result["reasoning_content"] = reasoning_content

                        # 添加用量信息
                        if hasattr(response, 'usage'):
                            result["usage"] = response.usage

                        return result

            # 如果没有找到内容，返回错误
            return {
                "success": False,
                "error": "API 响应格式异常，无法提取内容",
                "raw_response": str(response)
            }

        except Exception as e:
            return {
                "success": False,
                "error": f"响应解析失败: {str(e)}",
            }

    async def test_connection(self) -> Dict[str, Any]:
        """
        测试 API 连接（仅验证密钥有效性）

        Returns:
            Dict[str, Any]: 测试结果
        """
        try:
            # 验证 API 密钥是否有效
            if not self.api_key:
                return {
                    "success": False,
                    "message": "API 密钥未设置",
                }

            # 简单验证 API 密钥格式
            if not self.api_key.startswith("sk-"):
                return {
                    "success": False,
                    "message": "API 密钥格式无效（应以 'sk-' 开头）",
                }

            return {
                "success": True,
                "message": "API 配置验证成功",
                "model": self.model,
            }

        except Exception as e:
            return {
                "success": False,
                "message": f"API 配置验证失败: {str(e)}",
            }

    async def generate_image(
        self,
        prompt: str,
        size: str = '1328*1328',
        watermark: bool = False,
        prompt_extend: bool = True,
        negative_prompt: str = ''
    ) -> Dict[str, Any]:
        """
        生成图片

        Args:
            prompt: 图片描述文本
            size: 图片尺寸，支持 '1024*1024', '1328*1328', '720*1280', '1280*720'
            watermark: 是否添加水印
            prompt_extend: 是否扩展提示词
            negative_prompt: 负面提示词

        Returns:
            Dict[str, Any]: API 响应结果，包含 url、revised_prompt 等信息

        Raises:
            Exception: API 调用失败或参数无效
        """
        try:
            # 验证尺寸参数
            if size not in IMAGE_SIZES:
                valid_sizes = ', '.join(IMAGE_SIZES.keys())
                raise ValueError(f"无效的尺寸: {size}。支持的尺寸: {valid_sizes}")

            if not prompt or not isinstance(prompt, str):
                raise ValueError("提示词不能为空")

            # 构建消息内容
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"text": prompt}
                    ]
                }
            ]

            # 调用 MultiModalConversation API 生成图片
            response = MultiModalConversation.call(
                model="qwen-image-plus",
                messages=messages,
                result_format='message',
                stream=False,
                watermark=watermark,
                prompt_extend=prompt_extend,
                negative_prompt=negative_prompt,
                size=size
            )

            return self._parse_generate_response(response)

        except ValueError as e:
            raise ValueError(f"参数错误: {str(e)}")
        except Exception as e:
            raise Exception(f"图片生成失败: {str(e)}")

    async def edit_image(
        self,
        images: list[dict[str, str]],
        prompt: str,
        n: int = 1,
        negative_prompt: str = " ",
        prompt_extend: bool = True,
        watermark: bool = False,
        size: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        编辑图片

        Args:
            images: 图片列表，每个元素包含 'image' (base64/url) 和可选的 'text' (编辑指令)
                    例如: [{'image': '...'}, {'image': '...'}, {'text': '...'}]
                    或者直接传入 prompt 参数作为指令
            prompt: 编辑指令
            n: 生成数量 (1-6)
            negative_prompt: 负面提示词
            prompt_extend: 是否扩展提示词
            watermark: 是否添加水印
            size: 输出尺寸 (仅当 n=1 时有效)

        Returns:
            Dict[str, Any]: API 响应结果
        """
        try:
            # 构建消息内容
            content = []
            for img in images:
                if 'image' in img:
                    content.append({'image': img['image']})
            
            # 添加编辑指令
            content.append({'text': prompt})

            messages = [
                {
                    "role": "user",
                    "content": content
                }
            ]

            # 构建参数
            kwargs = {
                "model": "qwen-image-edit-plus",
                "messages": messages,
                "stream": False,
                "n": n,
                "negative_prompt": negative_prompt,
                "prompt_extend": prompt_extend,
                "watermark": watermark
            }

            if n == 1 and size:
                kwargs["size"] = size

            # 调用 API
            response = MultiModalConversation.call(**kwargs)

            return self._parse_generate_response(response)

        except Exception as e:
            raise Exception(f"图片编辑失败: {str(e)}")

    def _parse_generate_response(self, response: Any) -> Dict[str, Any]:
        """
        解析图片生成 API 响应

        Args:
            response: DashScope API 响应对象

        Returns:
            Dict[str, Any]: 解析后的响应数据，包含 url、revised_prompt 等字段
        """
        try:
            if not response:
                return {
                    "success": False,
                    "error": "API 响应为空"
                }

            # 检查响应状态
            if hasattr(response, 'status_code') and response.status_code != 200:
                error_msg = getattr(response, 'message', f"HTTP {response.status_code}")
                return {
                    "success": False,
                    "error": f"API 错误: {error_msg}",
                    "code": getattr(response, 'code', None)
                }

            image_url = None
            revised_prompt = None

            # 尝试从 response.output.choices[0].message.content 中提取图片 URL
            try:
                if hasattr(response, 'output') and response.output:
                    output = response.output
                    if hasattr(output, 'choices') and output.choices and len(output.choices) > 0:
                        choice = output.choices[0]

                        # 尝试从 message 中提取内容
                        if hasattr(choice, 'message'):
                            message = choice.message

                            # 提取 content（可能是列表或单个值）
                            content = getattr(message, 'content', [])
                            if isinstance(content, list) and len(content) > 0:
                                for item in content:
                                    if isinstance(item, dict):
                                        # 尝试多个可能的键名
                                        if 'url' in item:
                                            image_url = item.get('url')
                                        elif 'image' in item:
                                            image_url = item.get('image')
                                    elif hasattr(item, 'url'):
                                        image_url = getattr(item, 'url', None)
                                    elif hasattr(item, 'image'):
                                        image_url = getattr(item, 'image', None)

                                    if image_url:
                                        break

                            # 尝试从 message 的不同位置提取 revised_prompt
                            if hasattr(message, 'revised_prompt'):
                                revised_prompt = getattr(message, 'revised_prompt', None)

                        # 尝试从 choice 中提取 revised_prompt（如果不在 message 中）
                        if not revised_prompt and hasattr(choice, 'revised_prompt'):
                            revised_prompt = getattr(choice, 'revised_prompt', None)
            except Exception as e:
                # 记录但不中断，继续尝试其他方式
                pass

            # 如果从标准路径没有获取到 URL，尝试其他可能的路径
            if not image_url:
                try:
                    # 尝试直接从 response 获取
                    if hasattr(response, 'data'):
                        data = response.data
                        if isinstance(data, dict) and 'url' in data:
                            image_url = data['url']
                except:
                    pass

            # 成功提取了图片 URL
            if image_url:
                result = {
                    "success": True,
                    "url": image_url,
                    "model": "qwen-image-plus",
                }

                # 添加 revised_prompt（如果有）
                if revised_prompt:
                    result["revised_prompt"] = revised_prompt

                # 添加用量信息
                try:
                    if hasattr(response, 'usage'):
                        result["usage"] = response.usage
                except:
                    pass

                return result

            # 如果没有找到图片 URL
            return {
                "success": False,
                "error": "API 响应中无法找到图片 URL",
                "raw_response": str(response)
            }

        except Exception as e:
            return {
                "success": False,
                "error": f"响应解析异常: {str(e)}",
            }

    def get_model_info(self) -> Dict[str, str]:
        """
        获取模型信息

        Returns:
            Dict[str, str]: 模型信息
        """
        return {
            "model": self.model,
            "provider": "DashScope",
            "sdk": "dashscope-python",
            "capabilities": "图像理解、视觉问答、内容分析、思考推理、图像生成",
        }


# 创建全局客户端实例
_client: Optional[DashScopeClient] = None


def get_client() -> DashScopeClient:
    """获取全局客户端实例"""
    global _client
    if _client is None:
        _client = DashScopeClient()
    return _client


def set_client(client: DashScopeClient) -> None:
    """设置全局客户端实例"""
    global _client
    _client = client
