"""
MCP 服务器主逻辑

实现 MCP 协议的服务器，提供图片分析工具。
"""

import asyncio
import sys
import json
import logging
from typing import Any, Dict, List, Optional

try:
    from mcp.server import Server
    from mcp.server.models import InitializationOptions
    from mcp.server.stdio import stdio_server
    from mcp.types import Tool, TextContent, ImageContent
except ImportError:
    raise ImportError("MCP is required. Install with: pip install mcp>=1.0.0")

from .image_handler import image_handler
from .api_client import get_client, DashScopeClient, IMAGE_SIZES
from .image_downloader import get_downloader

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("image-mcp-server")


class ImageMCPServer:
    """Image MCP 服务器类"""

    def __init__(self):
        """初始化 MCP 服务器"""
        self.server = Server("image-mcp-server")
        self._setup_handlers()
        self._client: Optional[DashScopeClient] = None

    def _setup_handlers(self):
        """设置 MCP 处理器"""

        @self.server.list_tools()
        async def list_tools() -> List[Tool]:
            """列出可用工具"""
            # 构建尺寸选项的枚举
            size_enum = list(IMAGE_SIZES.keys())

            return [
                Tool(
                    name="analyze_image",
                    description="分析图片内容，支持本地文件、HTTP URL 和剪贴板图片。基于通义千问3-VL-Plus模型进行视觉理解和思考推理。",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "image": {
                                "type": "string",
                                "description": "图片来源：本地文件路径、HTTP URL 或 'clipboard' 表示剪贴板图片"
                            },
                            "question": {
                                "type": "string",
                                "description": "关于图片的问题（可选）",
                                "default": "请分析这张图片的内容"
                            },
                            "enable_thinking": {
                                "type": "boolean",
                                "description": "是否启用思考功能，让模型显示推理过程（可选）",
                                "default": True
                            },
                            "thinking_budget": {
                                "type": "integer",
                                "description": "思考过程最大 Token 数，范围 0-81920（可选）",
                                "default": 81920
                            }
                        },
                        "required": ["image"]
                    }
                ),
                Tool(
                    name="generate_image",
                    description="使用通义千问图像生成模型 (qwen-image-plus) 根据文本描述生成图片。支持自定义尺寸、水印和负面提示词等参数。",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "prompt": {
                                "type": "string",
                                "description": "图片描述文本"
                            },
                            "size": {
                                "type": "string",
                                "description": f"图片尺寸，支持: {', '.join(size_enum)}（可选）",
                                "enum": size_enum,
                                "default": "1328*1328"
                            },
                            "save_to_file": {
                                "type": "boolean",
                                "description": "是否下载保存到本地（可选）",
                                "default": False
                            },
                            "save_path": {
                                "type": "string",
                                "description": "保存路径（仅当 save_to_file=True 时有效，可选）"
                            },
                            "watermark": {
                                "type": "boolean",
                                "description": "是否添加水印（可选）",
                                "default": False
                            },
                            "prompt_extend": {
                                "type": "boolean",
                                "description": "是否扩展提示词（可选）",
                                "default": True
                            },
                            "negative_prompt": {
                                "type": "string",
                                "description": "负面提示词，描述不想要的内容（可选）",
                                "default": ""
                            }
                        },
                        "required": ["prompt"]
                    }
                ),
                Tool(
                    name="edit_image",
                    description="使用通义千问图像编辑模型 (qwen-image-edit-plus) 编辑图片。支持多图输入（最多3张），可精确修改文字、增删物体、迁移风格等。",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "images": {
                                "type": "array",
                                "items": {
                                    "type": "string"
                                },
                                "description": "输入图片列表（1-3张），支持文件路径、HTTP URL 或 'clipboard'。多图编辑时，请在 prompt 中用'图1'、'图2'等指代。",
                                "minItems": 1,
                                "maxItems": 3
                            },
                            "prompt": {
                                "type": "string",
                                "description": "编辑指令，描述期望的编辑效果。多图编辑时需明确指代（如'图1中的...'）。"
                            },
                            "n": {
                                "type": "integer",
                                "description": "生成数量 (1-6)，默认 1",
                                "default": 1,
                                "minimum": 1,
                                "maximum": 6
                            },
                            "size": {
                                "type": "string",
                                "description": f"输出尺寸 (仅当 n=1 时有效)，支持: {', '.join(size_enum)}（可选）",
                                "enum": size_enum
                            },
                            "save_to_file": {
                                "type": "boolean",
                                "description": "是否下载保存到本地（可选）",
                                "default": False
                            },
                            "save_path": {
                                "type": "string",
                                "description": "保存路径（仅当 save_to_file=True 时有效，可选）"
                            },
                            "watermark": {
                                "type": "boolean",
                                "description": "是否添加水印（可选）",
                                "default": False
                            },
                            "prompt_extend": {
                                "type": "boolean",
                                "description": "是否扩展提示词（可选）",
                                "default": True
                            },
                            "negative_prompt": {
                                "type": "string",
                                "description": "负面提示词（可选）",
                                "default": " "
                            }
                        },
                        "required": ["images", "prompt"]
                    }
                )
            ]

        @self.server.call_tool()
        async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
            """调用工具"""
            if name == "analyze_image":
                return await self._handle_analyze_image(arguments)
            elif name == "generate_image":
                return await self._handle_generate_image(arguments)
            elif name == "edit_image":
                return await self._handle_edit_image(arguments)
            else:
                raise ValueError(f"未知工具: {name}")

        @self.server.list_prompts()
        async def list_prompts():
            """列出可用提示（空列表，因为只有工具）"""
            return []

        @self.server.get_prompt()
        async def get_prompt(name: str, arguments: Optional[Dict[str, str]] = None):
            """获取提示（不支持）"""
            raise ValueError("不支持提示功能")

    async def _handle_analyze_image(self, arguments: Dict[str, Any]) -> List[TextContent]:
        """
        处理图片分析请求

        Args:
            arguments: 工具参数

        Returns:
            List[TextContent]: 分析结果
        """
        try:
            # 解析参数
            image_input = arguments.get("image", "")
            question = arguments.get("question", "请分析这张图片的内容")
            enable_thinking = arguments.get("enable_thinking", True)
            thinking_budget = arguments.get("thinking_budget", 81920)

            if not image_input:
                return [TextContent(
                    type="text",
                    text="错误: 缺少图片参数。请提供图片文件路径、HTTP URL 或使用 'clipboard' 来分析剪贴板图片。"
                )]

            logger.info(f"开始分析图片: {image_input}")
            logger.info(f"思考功能: {enable_thinking}, Token 预算: {thinking_budget}")

            # 处理图片
            try:
                base64_data, mime_type = image_handler.process_image_input(image_input)
                logger.info(f"图片处理成功，格式: {mime_type}")
            except Exception as e:
                error_msg = f"图片处理失败: {str(e)}"
                logger.error(error_msg)
                return [TextContent(type="text", text=error_msg)]

            # 调用 API 分析图片
            try:
                client = get_client()
                result = await client.analyze_image(
                    base64_data,
                    mime_type,
                    question,
                    enable_thinking=enable_thinking,
                    thinking_budget=thinking_budget
                )

                if result.get("success", False):
                    # 构建响应内容
                    response_text = ""

                    # 如果有思考过程，添加到响应
                    if "reasoning_content" in result and result["reasoning_content"]:
                        response_text += "【思考过程】\n"
                        response_text += result["reasoning_content"]
                        response_text += "\n\n【分析结果】\n"

                    # 添加分析结果
                    analysis = result.get("analysis", "分析完成但没有返回内容")
                    response_text += analysis

                    logger.info("图片分析成功")
                    return [TextContent(type="text", text=response_text)]
                else:
                    error_msg = result.get("error", "图片分析失败")
                    logger.error(f"API分析失败: {error_msg}")
                    return [TextContent(type="text", text=f"分析失败: {error_msg}")]

            except Exception as e:
                error_msg = f"API调用失败: {str(e)}"
                logger.error(error_msg)
                return [TextContent(type="text", text=error_msg)]

        except Exception as e:
            error_msg = f"处理请求时发生错误: {str(e)}"
            logger.error(error_msg)
            return [TextContent(type="text", text=error_msg)]

    async def _handle_generate_image(self, arguments: Dict[str, Any]) -> List[TextContent]:
        """
        处理图片生成请求

        Args:
            arguments: 工具参数

        Returns:
            List[TextContent]: 生成结果
        """
        try:
            # 解析参数
            prompt = arguments.get("prompt", "")
            size = arguments.get("size", "1024*1024")
            save_to_file = arguments.get("save_to_file", False)
            save_path = arguments.get("save_path")
            watermark = arguments.get("watermark", False)
            prompt_extend = arguments.get("prompt_extend", True)
            negative_prompt = arguments.get("negative_prompt", "")

            if not prompt:
                return [TextContent(
                    type="text",
                    text="错误: 缺少 prompt 参数。请提供图片描述文本。"
                )]

            logger.info(f"开始生成图片，提示词: {prompt[:50]}...")
            logger.info(f"参数: 尺寸={size}, 水印={watermark}, 扩展提示词={prompt_extend}")

            # 调用 API 生成图片
            try:
                client = get_client()
                result = await client.generate_image(
                    prompt=prompt,
                    size=size,
                    watermark=watermark,
                    prompt_extend=prompt_extend,
                    negative_prompt=negative_prompt
                )

                if result.get("success", False):
                    image_url = result.get("url", "")
                    revised_prompt = result.get("revised_prompt")

                    # 构建响应内容
                    response_text = "✓ 图片生成成功！\n\n"
                    response_text += "📸 生成的图片:\n"
                    response_text += f"- URL: {image_url}\n"
                    response_text += f"- 尺寸: {size}\n"

                    if revised_prompt:
                        response_text += f"- 优化后提示词: {revised_prompt}\n"

                    # 如果请求保存到本地
                    if save_to_file:
                        try:
                            logger.info(f"正在下载图片到本地...")
                            downloader = get_downloader()
                            local_path = downloader.download_image(
                                image_url,
                                save_path=save_path
                            )
                            logger.info(f"图片已保存到: {local_path}")
                            response_text += f"\n💾 已保存到本地:\n"
                            response_text += f"- 路径: {local_path}\n"
                        except Exception as e:
                            logger.warning(f"保存图片失败: {str(e)}")
                            response_text += f"\n⚠️ 保存到本地失败: {str(e)}\n"

                    logger.info("图片生成成功")
                    return [TextContent(type="text", text=response_text)]
                else:
                    error_msg = result.get("error", "图片生成失败")
                    logger.error(f"API生成失败: {error_msg}")
                    return [TextContent(type="text", text=f"生成失败: {error_msg}")]

            except ValueError as e:
                error_msg = f"参数错误: {str(e)}"
                logger.error(error_msg)
                return [TextContent(type="text", text=error_msg)]
            except Exception as e:
                error_msg = f"API调用失败: {str(e)}"
                logger.error(error_msg)
                return [TextContent(type="text", text=error_msg)]

        except Exception as e:
            error_msg = f"处理请求时发生错误: {str(e)}"
            logger.error(error_msg)
            return [TextContent(type="text", text=error_msg)]

    async def _handle_edit_image(self, arguments: Dict[str, Any]) -> List[TextContent]:
        """
        处理图片编辑请求

        Args:
            arguments: 工具参数

        Returns:
            List[TextContent]: 编辑结果
        """
        try:
            # 解析参数
            images_input = arguments.get("images", [])
            prompt = arguments.get("prompt", "")
            n = arguments.get("n", 1)
            size = arguments.get("size")
            save_to_file = arguments.get("save_to_file", False)
            save_path = arguments.get("save_path")
            watermark = arguments.get("watermark", False)
            prompt_extend = arguments.get("prompt_extend", True)
            negative_prompt = arguments.get("negative_prompt", " ")

            if not images_input:
                return [TextContent(
                    type="text",
                    text="错误: 缺少 images 参数。请提供至少一张图片。"
                )]
            
            if not prompt:
                return [TextContent(
                    type="text",
                    text="错误: 缺少 prompt 参数。请提供编辑指令。"
                )]

            logger.info(f"开始编辑图片，输入图片数: {len(images_input)}, 提示词: {prompt[:50]}...")

            # 处理输入图片
            processed_images = []
            for i, img_input in enumerate(images_input):
                try:
                    base64_data, mime_type = image_handler.process_image_input(img_input)
                    # 转换为 data URI 格式
                    image_url = f"data:{mime_type};base64,{base64_data}"
                    processed_images.append({'image': image_url})
                    logger.info(f"图片 {i+1} 处理成功")
                except Exception as e:
                    error_msg = f"处理图片 {i+1} ({img_input}) 失败: {str(e)}"
                    logger.error(error_msg)
                    return [TextContent(type="text", text=error_msg)]

            # 调用 API 编辑图片
            try:
                client = get_client()
                result = await client.edit_image(
                    images=processed_images,
                    prompt=prompt,
                    n=n,
                    negative_prompt=negative_prompt,
                    prompt_extend=prompt_extend,
                    watermark=watermark,
                    size=size
                )

                if result.get("success", False):
                    image_url = result.get("url", "")
                    revised_prompt = result.get("revised_prompt")

                    # 构建响应内容
                    response_text = "✓ 图片编辑成功！\n\n"
                    response_text += "📸 生成的图片:\n"
                    
                    # 检查是否有多个结果 (虽然当前 api_client 只返回第一个 url，但为了兼容性...)
                    # api_client.edit_image 使用 _parse_generate_response，目前只提取第一个 URL
                    # 如果需要支持多图返回，可能需要修改 api_client._parse_generate_response 或在此处特殊处理
                    # 暂时假设 api_client 返回单个 url，或者我们需要修改 api_client 来支持多图返回
                    # 这里的 result['url'] 是单个 URL
                    
                    response_text += f"- URL: {image_url}\n"
                    
                    if n > 1:
                        response_text += "(注: 当前 API 客户端仅返回第一张生成的图片 URL，如需更多请联系开发者更新)\n"

                    if revised_prompt:
                        response_text += f"- 优化后提示词: {revised_prompt}\n"

                    # 如果请求保存到本地
                    if save_to_file:
                        try:
                            logger.info(f"正在下载图片到本地...")
                            downloader = get_downloader()
                            # 如果 save_path 是目录，自动生成文件名
                            local_path = downloader.download_image(
                                image_url,
                                save_path=save_path
                            )
                            logger.info(f"图片已保存到: {local_path}")
                            response_text += f"\n💾 已保存到本地:\n"
                            response_text += f"- 路径: {local_path}\n"
                        except Exception as e:
                            logger.warning(f"保存图片失败: {str(e)}")
                            response_text += f"\n⚠️ 保存到本地失败: {str(e)}\n"

                    logger.info("图片编辑成功")
                    return [TextContent(type="text", text=response_text)]
                else:
                    error_msg = result.get("error", "图片编辑失败")
                    logger.error(f"API编辑失败: {error_msg}")
                    return [TextContent(type="text", text=f"编辑失败: {error_msg}")]

            except Exception as e:
                error_msg = f"API调用失败: {str(e)}"
                logger.error(error_msg)
                return [TextContent(type="text", text=error_msg)]

        except Exception as e:
            error_msg = f"处理请求时发生错误: {str(e)}"
            logger.error(error_msg)
            return [TextContent(type="text", text=error_msg)]

    async def run(self):
        """运行 MCP 服务器"""
        # 验证 API 配置
        try:
            client = get_client()
            test_result = await client.test_connection()
            if test_result.get("success"):
                logger.info(f"API配置验证成功: {test_result.get('message')}")
            else:
                logger.warning(f"API配置验证失败: {test_result.get('message')}")
        except Exception as e:
            logger.error(f"API配置验证异常: {str(e)}")

        # 启动服务器
        logger.info("Image MCP Server 启动...")
        async with stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream,
                write_stream,
                InitializationOptions(
                    server_name="image-mcp-server",
                    server_version="0.3.1",
                    capabilities={
                        "tools": {}
                    }
                )
            )


# 全局服务器实例
_server_instance: Optional[ImageMCPServer] = None


def get_server() -> ImageMCPServer:
    """获取服务器实例"""
    global _server_instance
    if _server_instance is None:
        _server_instance = ImageMCPServer()
    return _server_instance


def main():
    """主函数 - 同步入口点"""
    try:
        server = get_server()
        asyncio.run(server.run())
    except KeyboardInterrupt:
        logger.info("收到中断信号，正在关闭服务器...")
    except Exception as e:
        logger.error(f"服务器运行错误: {str(e)}")
        logger.error(f"错误类型: {type(e).__name__}")
        import traceback
        logger.error(f"详细错误信息:\n{traceback.format_exc()}")
        sys.exit(1)


if __name__ == "__main__":
    main()