# -*- coding: utf-8 -*-
# filename: server.py
# @Time    : 2025-01-20
# @Author  : Claude Code
# @Email   : noreply@anthropic.com
# @Software: PyCharm
"""MCP Server for TF Tools - Image generation and understanding tools"""
import logging
import os
from pathlib import Path
from typing import Any, Literal

import anyio
from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel

from tf_tools_mcp.tools.img_gen_tools.text_to_image import TextToImageTool
from tf_tools_mcp.tools.img_gen_tools.image_to_image import ImageToImageTool
from tf_tools_mcp.tools.file_preview_tools.image_understanding import (
    ImageUnderstandingTool,
)
from tf_tools_mcp.tools.file_preview_tools.file_preview import FilePreviewTool
from tf_tools_mcp.tools.file_preview_tools.file_preview_types import FilePreviewParams
from tf_tools_mcp.tools.img_gen_tools.types import (
    TextToImageMode,
    ImageToImageMode,
    ImageGenerationResult,
)

# ============== Literal 类型定义，用于枚举参数 ==============

# 支持的图像生成模型
GenerationModel = Literal[
    "google/gemini-3-pro-image-preview",
    "bytedance-seed/seedream-4.5",
    "openai/gpt-5-image",
    "openai/gpt-5-image-mini",
    "black-forest-labs/flux.2-max",
    "black-forest-labs/flux.2-pro",
]

# 图像质量选项
ImageQuality = Literal["low", "medium", "high"]

# 图像布局选项
ImageLayout = Literal["square", "landscape", "portrait"]

# 图像理解模型
UnderstandingModel = Literal["glm-4.6v", "glm-4.6v-flash", "glm-4.6v-flashx"]

# 图像理解提示词模板
PromptTemplate = Literal[
    "describe",
    "ocr",
    "detect_objects",
    "analyze_chart",
    "describe_pdf",
    "describe_ppt",
    "custom",
]

logger = logging.getLogger(__name__)

# Initialize FastMCP server
app = FastMCP(
    name="tf-tools-mcp",
    instructions="Image generation and understanding tools using OpenRouter and Zhipu AI",
)

# Global tool instances (initialized in lifespan)
_tools: dict[str, Any] = {}


@app.tool()
async def text_to_image(
    prompt: str,
    model: GenerationModel,
    mode: TextToImageMode = "custom",
    quality: ImageQuality = "medium",
    layout: ImageLayout = "square",
    save_to_local: bool = False,
    save_dir: str = "/tmp/generate_img/",
) -> ImageGenerationResult:
    """文本生成图片工具 - 使用 OpenRouter API 根据文本描述生成图片

    支持多个主流 AI 图像生成模型，包括 Gemini 3 Pro、Seedream 4.5、GPT-5 Image、FLUX.2 等。
    可选使用 GLM-4.6V 进行 prompt 增强，提升生成质量。

    入参说明：
        prompt (str, 必填): 图像生成的文本描述，例如"一只可爱的猫咪在花园里玩耍"
        model (枚举, 必填): 使用的图像生成模型
            - google/gemini-3-pro-image-preview: Google Gemini 3 Pro 图像预览版
            - bytedance-seed/seedream-4.5: 字节跳动 Seedream 4.5
            - openai/gpt-5-image: OpenAI GPT-5 图像生成
            - openai/gpt-5-image-mini: OpenAI GPT-5 迷你版
            - black-forest-labs/flux.2-max: FLUX.2 Max 专业版
            - black-forest-labs/flux.2-pro: FLUX.2 Pro 标准版
        mode (枚举, 可选, 默认 custom): prompt 增强模式
            - technical_architecture: 技术架构图 - 生成系统架构图、数据流图等
            - design_mockup: 设计稿 - 生成 UI 设计稿、线框图等
            - promotional_material: 宣传物料 - 生成海报、banner 等
            - office_chart: 通用办公图表 - 生成信息图表、数据可视化等
            - custom: 自定义 - 不进行 prompt 增强
        quality (枚举, 可选, 默认 medium): 图像质量
            - low: 低质量（生成速度快）
            - medium: 中等质量（平衡速度和质量）
            - high: 高质量（生成速度慢）
        layout (枚举, 可选, 默认 square): 图像布局
            - square: 正方形 (1:1)
            - landscape: 横向 (16:9)
            - portrait: 纵向 (9:16)
        save_to_local (bool, 可选, 默认 False): 是否保存到本地
            - False: 返回 base64 编码的图片数据（默认）
            - True: 将图片保存到本地文件，返回文件路径
            - 设置为 True 可避免返回大量 base64 数据
        save_dir (str, 可选, 默认 /tmp/generate_img/): 本地保存目录
            - 仅当 save_to_local=True 时生效
            - 目录不存在时会自动创建
            - 支持绝对路径和相对路径
            - 示例: "/tmp/images/", "./output/", "~/generated_images/"

    出参说明：
        返回 ImageGenerationResult 对象，包含：

        当 save_to_local=False 时：
        - base64_data: base64 编码的图像数据（data:image/png;base64,... 或 data:image/jpeg;base64,... 格式）
        - file_path: None

        当 save_to_local=True 时：
        - base64_data: None
        - file_path: 图片文件的绝对路径（例如: /tmp/generate_img/google_gemini-3-pro_20250109_163000.png）

        其他字段（无论 save_to_local 取值）：
        - width: 图像宽度
        - height: 图像高度
        - file_size: 文件大小（字节）
        - format: 图像格式（PNG、JPEG 等）
        - generated_at: 生成时间
        - model: 使用的模型
        - quality: 图像质量
        - layout: 图像布局
        - original_prompt: 用户输入的原始 prompt（始终有值）
        - enhanced_prompt: 增强后的 prompt（仅在使用 prompt 增强时有值，否则为 null）
    """
    tool_key = f"text_to_image_{model}"
    if tool_key not in _tools:
        storage_dir = os.environ.get("IMAGE_STORAGE_DIR")  # None if not set
        _tools[tool_key] = TextToImageTool(
            model=model,
            storage_dir=storage_dir,
            enable_prompt_enhancement=True,
        )

    tool = _tools[tool_key]

    result = await anyio.to_thread.run_sync(
        lambda: tool.run(
            prompt=prompt,
            mode=mode,
            quality=quality,
            layout=layout,
            save_to_local=save_to_local,
            save_dir=save_dir,
        )
    )

    return result


@app.tool()
async def image_to_image(
    prompt: str,
    reference_images: list[str],
    model: GenerationModel,
    mode: ImageToImageMode = "custom",
    quality: ImageQuality = "medium",
    layout: ImageLayout = "square",
    save_to_local: bool = False,
    save_dir: str = "/tmp/generate_img/",
) -> ImageGenerationResult:
    """图片生成图片工具 - 使用参考图片和文本描述生成新图片

    支持多个主流 AI 图像生成模型，可以根据参考图片的风格和内容，结合文本描述生成新的图片。
    可选使用 GLM-4.6V 进行 prompt 增强。

    入参说明：
        prompt (str, 必填): 图像生成的文本描述，例如"将背景改为日落景色"
        reference_images (list[str], 必填): 参考图片列表，支持 1-8 张图片
            - 本地文件路径: /path/to/image.jpg
            - 网络 URL: https://example.com/image.jpg
            - 最少 1 张，最多 8 张图片
        model (枚举, 必填): 使用的图像生成模型
            - google/gemini-3-pro-image-preview: Google Gemini 3 Pro 图像预览版
            - bytedance-seed/seedream-4.5: 字节跳动 Seedream 4.5
            - openai/gpt-5-image: OpenAI GPT-5 图像生成
            - openai/gpt-5-image-mini: OpenAI GPT-5 迷你版
            - black-forest-labs/flux.2-max: FLUX.2 Max 专业版
            - black-forest-labs/flux.2-pro: FLUX.2 Pro 标准版
        mode (枚举, 可选, 默认 custom): prompt 增强模式
            - content_adjustment: 内容调整 - 调整亮度、对比度、色彩平衡等
            - element_replace: 元素替换 - 替换图中的人物、物体、背景等
            - style_transfer: 风格转换 - 转为油画、卡通、水彩等风格
            - text_modification: 文字修改 - 修改图片内的文字内容
            - custom: 自定义 - 不进行 prompt 增强
        quality (枚举, 可选, 默认 medium): 图像质量
            - low: 低质量（生成速度快）
            - medium: 中等质量（平衡速度和质量）
            - high: 高质量（生成速度慢）
        layout (枚举, 可选, 默认 square): 图像布局
            - square: 正方形 (1:1)
            - landscape: 横向 (16:9)
            - portrait: 纵向 (9:16)
        save_to_local (bool, 可选, 默认 False): 是否保存到本地
            - False: 返回 base64 编码的图片数据（默认）
            - True: 将图片保存到本地文件，返回文件路径
            - 设置为 True 可避免返回大量 base64 数据
        save_dir (str, 可选, 默认 /tmp/generate_img/): 本地保存目录
            - 仅当 save_to_local=True 时生效
            - 目录不存在时会自动创建
            - 支持绝对路径和相对路径
            - 示例: "/tmp/images/", "./output/", "~/generated_images/"

    出参说明：
        返回 ImageGenerationResult 对象，包含：

        当 save_to_local=False 时：
        - base64_data: base64 编码的图像数据（data:image/png;base64,... 或 data:image/jpeg;base64,... 格式）
        - file_path: None

        当 save_to_local=True 时：
        - base64_data: None
        - file_path: 图片文件的绝对路径（例如: /tmp/generate_img/google_gemini-3-pro_img2img_20250109_163000.png）

        其他字段（无论 save_to_local 取值）：
        - width: 图像宽度
        - height: 图像高度
        - file_size: 文件大小（字节）
        - format: 图像格式（PNG、JPEG 等）
        - generated_at: 生成时间
        - model: 使用的模型
        - quality: 图像质量
        - layout: 图像布局
        - original_prompt: 用户输入的原始 prompt（始终有值）
        - enhanced_prompt: 增强后的 prompt（仅在使用 prompt 增强时有值，否则为 null）
    """
    # 验证图片数量
    if not reference_images or len(reference_images) == 0:
        raise ValueError("reference_images 不能为空，至少需要 1 张图片")
    if len(reference_images) > 8:
        raise ValueError(f"reference_images 最多支持 8 张图片，当前传入了 {len(reference_images)} 张")

    tool_key = f"image_to_image_{model}"
    if tool_key not in _tools:
        storage_dir = os.environ.get("IMAGE_STORAGE_DIR")  # None if not set
        _tools[tool_key] = ImageToImageTool(
            model=model,
            storage_dir=storage_dir,
            enable_prompt_enhancement=True,
        )

    tool = _tools[tool_key]

    result = await anyio.to_thread.run_sync(
        lambda: tool.run(
            prompt=prompt,
            reference_images=reference_images,
            mode=mode,
            quality=quality,
            layout=layout,
            save_to_local=save_to_local,
            save_dir=save_dir,
        )
    )

    return result


@app.tool()
async def image_understanding(
    images: list[str],
    prompt_template: PromptTemplate = "describe",
    custom_prompt: str | None = None,
    model: UnderstandingModel = "glm-4.6v-flashx",
) -> str:
    """图片理解工具 - 使用智谱 GLM-4.6V 视觉模型分析和理解图片内容

    支持单图或批量图片分析，提供多种预置分析模板：图片描述、OCR 文字提取、对象检测、图表分析等。

    入参说明：
        images (list[str], 必填): 图片路径列表，支持本地文件路径或网络 URL
            - 本地文件: /path/to/image.jpg
            - 网络 URL: https://example.com/image.jpg
        prompt_template (枚举, 可选, 默认 describe): 分析模板
            - describe: 全面描述 - 分析图片的对象、场景、颜色、风格、氛围等
            - ocr: 文字提取 - 提取图片中的所有文字内容，保持格式
            - detect_objects: 对象检测 - 识别图片中的主要物体、次要物体、背景等
            - analyze_chart: 图表分析 - 分析图表类型、数据趋势、关键数据点等
            - describe_pdf: PDF 描述 - 专门用于 PDF 文档页面转换的图片
            - describe_ppt: PPT 描述 - 专门用于 PPT 幻灯片转换的图片
            - custom: 自定义 - 使用自定义 prompt
        custom_prompt (str, 可选): 自定义分析提示词
            - 仅当 prompt_template 为 "custom" 时必填
            - 例如: "请分析这张图片中的建筑物风格和历史年代"
        model (枚举, 可选, 默认 glm-4.6v-flashx): 使用的 GLM 视觉模型
            - glm-4.6v: 高性能模型 - 适合复杂视觉理解、高精度 OCR
            - glm-4.6v-flash: 免费快速模型 - 适合通用视觉分析
            - glm-4.6v-flashx: 超快速模型 - 适合实时分析

    出参说明：
        返回格式化的结果字符串，包含：
        - 图片分析结果（详细描述或提取的信息）
        - 使用的模型和模板
        - token 消耗统计（输入/输出/总计）
    """
    tool_key = f"image_understanding_{model}"
    try:
        if tool_key not in _tools:
            _tools[tool_key] = ImageUnderstandingTool(model=model)

        tool = _tools[tool_key]

        result = await anyio.to_thread.run_sync(
            lambda: tool.run(
                images=images,
                prompt_template=prompt_template,
                custom_prompt=custom_prompt,
            )
        )

        return tool.format_output(result)
    except ValueError as e:
        # 返回友好的错误消息
        return str(e)


@app.tool()
async def file_preview(
    file_paths: list[str],
    options: dict[str, Any] | None = None,
) -> str:
    """文件预览工具 - 支持多种文件格式的预览和内容摘要生成

    支持图片、PDF、Word、Excel、PPT 等多种格式，自动提取元数据并使用 GLM-4.6V 生成内容摘要。
    支持本地文件和远程 URL。

    入参说明：
        file_paths (list[str], 必填): 文件路径列表，支持本地文件路径或网络 URL
            - 本地文件: /path/to/document.pdf
            - 网络 URL: https://example.com/document.pdf
        options (dict, 可选): 预览选项配置（所有字段均为可选，未提供时使用环境变量或默认值）
            - max_file_size (int): 最大文件大小（字节），默认 100MB
            - max_download_size (int): 最大下载大小（字节），默认 50MB
            - download_timeout (int): 下载超时时间（秒），默认 30
            - max_pages (int): 最大页数（PDF/PPT），默认 3
            - max_chars (int): 最大字符数（文档内容），默认 1000
            - max_preview_slides (int): 最大预览幻灯片数（PPT），默认 3
            - enable_cache (bool): 是否启用缓存，默认 true
            - max_workers (int): 最大并发工作线程数，默认 None（自动）

    支持的文件格式：
        - 图片: JPG, PNG, GIF, BMP, WebP 等
        - 文档: PDF, DOCX, DOC, TXT, MD 等
        - 表格: XLSX, XLS, CSV 等
        - 演示: PPTX, PPT 等

    出参说明：
        返回格式化的结果字符串，每个文件包含：
        - 文件基本信息（名称、类型、大小）
        - 元数据（创建时间、修改时间、作者等）
        - 内容摘要（自动生成）
        - 处理信息（耗时、使用的配置等）
    """
    options_dict = options or {}

    # Use hardcoded default values
    default_max_file_size = 100 * 1024 * 1024  # 100MB
    default_max_download_size = 50 * 1024 * 1024  # 50MB
    default_download_timeout = 30
    default_max_pages = 3
    default_max_chars = 1000
    default_max_preview_slides = 3
    default_enable_cache = True
    default_max_workers = None

    # Build options from options_dict
    from tf_tools_mcp.tools.file_preview_tools.file_preview_types import (
        FilePreviewOptions,
    )

    options_kwargs = {
        "max_file_size": options_dict.get("max_file_size", default_max_file_size),
        "max_download_size": options_dict.get(
            "max_download_size", default_max_download_size
        ),
        "download_timeout": options_dict.get(
            "download_timeout", default_download_timeout
        ),
        "max_pages": options_dict.get("max_pages", default_max_pages),
        "max_chars": options_dict.get("max_chars", default_max_chars),
        "max_preview_slides": options_dict.get(
            "max_preview_slides", default_max_preview_slides
        ),
        "enable_cache": options_dict.get("enable_cache", default_enable_cache),
    }

    # Handle max_workers separately (can be None)
    if "max_workers" in options_dict:
        options_kwargs["max_workers"] = options_dict["max_workers"]
    elif default_max_workers is not None:
        options_kwargs["max_workers"] = int(default_max_workers)

    # Create options with validation
    options = FilePreviewOptions(**options_kwargs)

    # Build params
    params = FilePreviewParams(file_paths=file_paths, options=options)

    # Get or create tool instance with default options
    tool_key = "file_preview"
    if tool_key not in _tools:
        # Create tool with default options
        default_options = FilePreviewOptions(
            max_file_size=default_max_file_size,
            max_download_size=default_max_download_size,
            download_timeout=default_download_timeout,
            max_pages=default_max_pages,
            max_chars=default_max_chars,
            max_preview_slides=default_max_preview_slides,
            enable_cache=default_enable_cache,
            max_workers=default_max_workers,
        )
        _tools[tool_key] = FilePreviewTool(default_options=default_options)

    tool = _tools[tool_key]

    # Run preview
    result = await tool.run(params)

    return FilePreviewTool.format_output(result)


def run_server():
    """Synchronous entry point for MCP server"""
    import logging

    logging.basicConfig(level=logging.INFO)

    # Create storage directory if IMAGE_STORAGE_DIR is configured
    storage_dir = os.environ.get("IMAGE_STORAGE_DIR")
    if storage_dir:
        Path(storage_dir).mkdir(parents=True, exist_ok=True)

    # Run the MCP server with stdio transport
    app.run(transport="stdio")


if __name__ == "__main__":
    run_server()
