import base64
import mimetypes
import os
import sys
from pathlib import Path
from typing import Annotated

import httpx
from openai import AsyncOpenAI
from pydantic import Field
from fastmcp import FastMCP

mcp = FastMCP(name="image-mcp")


async def resolve_image_input(
    image_input: str,
    force_base64: bool = False,
    max_size_mb: float = 20,
) -> dict:
    """Resolve image input to a standardized format.

    Args:
        image_input: HTTP URL, data URI, or local file path.
        force_base64: If True, always return base64 data (for OCR API).
        max_size_mb: Maximum file size in MB.

    Returns:
        dict with keys depending on type:
        - {"type": "url", "url": str} for HTTP URLs (when not force_base64)
        - {"type": "data_uri", "data_uri": str} for data URIs
        - {"type": "base64", "data": str} for raw base64 data
    """
    # HTTP URL
    if image_input.startswith(("http://", "https://")):
        if not force_base64:
            return {"type": "url", "url": image_input}
        try:
            async with httpx.AsyncClient(timeout=30) as client:
                resp = await client.get(image_input)
                resp.raise_for_status()
                content = resp.content
        except httpx.HTTPStatusError as e:
            raise ValueError(
                f"下载失败: {image_input} (HTTP {e.response.status_code})"
            ) from e
        except httpx.RequestError as e:
            raise ValueError(f"下载失败: {image_input} ({e})") from e
        max_bytes = int(max_size_mb * 1024 * 1024)
        if len(content) > max_bytes:
            raise ValueError(
                f"文件过大: {len(content) / 1024 / 1024:.1f}MB，最大支持 {max_size_mb}MB"
            )
        return {"type": "base64", "data": base64.b64encode(content).decode()}

    # Data URI
    if image_input.startswith("data:"):
        if "," not in image_input:
            raise ValueError("无效的 data URI 格式，缺少 base64 数据")
        if force_base64:
            _, encoded = image_input.split(",", 1)
            return {"type": "base64", "data": encoded}
        return {"type": "data_uri", "data_uri": image_input}

    # Local file path
    path = Path(image_input)
    if not path.exists():
        raise ValueError(
            f"文件不存在: {image_input}。"
            "如果 MCP server 与文件不在同一台机器，请使用 HTTP URL 或 base64 data URI"
        )

    file_size = path.stat().st_size
    max_bytes = int(max_size_mb * 1024 * 1024)
    if file_size > max_bytes:
        raise ValueError(
            f"文件过大: {file_size / 1024 / 1024:.1f}MB，最大支持 {max_size_mb}MB"
        )

    file_bytes = path.read_bytes()
    mime_type = mimetypes.guess_type(image_input)[0] or "application/octet-stream"

    if force_base64:
        return {"type": "base64", "data": base64.b64encode(file_bytes).decode()}

    b64 = base64.b64encode(file_bytes).decode()
    data_uri = f"data:{mime_type};base64,{b64}"
    return {"type": "data_uri", "data_uri": data_uri}


ANALYZE_IMAGE_FORMATS = {".jpg", ".jpeg", ".png", ".gif", ".webp"}
ANALYZE_IMAGE_MAX_SIZE_MB = 20


def get_openai_client() -> AsyncOpenAI:
    base_url = os.environ.get("LLM_BASE_URL")
    if not base_url:
        raise ValueError("环境变量 LLM_BASE_URL 未配置，请设置 Qwen VL API 地址")
    return AsyncOpenAI(
        base_url=base_url,
        api_key=os.environ.get("LLM_API_KEY", "no-key"),
    )


def _validate_image_format(image_input: str) -> None:
    """Validate file format for analyze_image. Only checks local paths."""
    if image_input.startswith(("http://", "https://", "data:")):
        return
    suffix = Path(image_input).suffix.lower()
    if suffix and suffix not in ANALYZE_IMAGE_FORMATS:
        raise ValueError(
            f"不支持的图片格式: {suffix}。"
            f"支持格式: {', '.join(sorted(ANALYZE_IMAGE_FORMATS))}"
        )


@mcp.tool
async def analyze_image(
    prompt: Annotated[str, Field(description="对图片的提问或分析要求")],
    image_url: Annotated[
        str,
        Field(
            description="图片来源，支持 HTTP/HTTPS URL、base64 data URI 或本地文件路径。"
            "支持格式：JPEG、PNG、GIF、WebP（最大 20MB）"
        ),
    ],
) -> str:
    """分析和理解图片内容。支持多种图片输入方式。"""
    _validate_image_format(image_url)
    resolved = await resolve_image_input(
        image_url, max_size_mb=ANALYZE_IMAGE_MAX_SIZE_MB
    )

    if resolved["type"] == "url":
        image_content_url = resolved["url"]
    else:
        image_content_url = resolved["data_uri"]

    client = get_openai_client()
    model = os.environ.get("LLM_MODEL", "Qwen3.5-27B")

    try:
        response = await client.chat.completions.create(
            model=model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {"url": image_content_url},
                        },
                    ],
                }
            ],
            stream=False,
            timeout=60,
        )
    except Exception as e:
        raise ValueError(f"LLM API 调用失败: {e}") from e

    if isinstance(response, str):
        return response
    return response.choices[0].message.content


OCR_FORMATS = {".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".webp", ".pdf"}
OCR_MAX_SIZE_MB = 50


def _validate_ocr_format(file_input: str) -> None:
    """Validate file format for OCR. Only checks local paths."""
    if file_input.startswith(("http://", "https://", "data:")):
        return
    suffix = Path(file_input).suffix.lower()
    if suffix and suffix not in OCR_FORMATS:
        raise ValueError(
            f"不支持的文件格式: {suffix}。"
            f"支持格式: {', '.join(sorted(OCR_FORMATS))}"
        )


def _detect_file_type(file_input: str) -> str:
    """Detect file type from input string."""
    if file_input.startswith("data:"):
        mime_part = file_input.split(";")[0].split("/")[-1]
        return mime_part
    suffix = Path(file_input.split("?")[0]).suffix.lstrip(".")
    return suffix.lower() if suffix else "png"


@mcp.tool
async def ocr(
    file_path: Annotated[
        str,
        Field(
            description="文件来源，支持 HTTP/HTTPS URL、base64 data URI 或本地文件路径。"
            "支持格式：PNG、JPG、BMP、TIFF、WebP、PDF（最大 50MB）"
        ),
    ],
) -> str:
    """OCR 文字识别，将图片或 PDF 转为结构化 Markdown 文本。"""
    _validate_ocr_format(file_path)
    resolved = await resolve_image_input(
        file_path, force_base64=True, max_size_mb=OCR_MAX_SIZE_MB
    )
    b64_data = resolved["data"]
    file_type = _detect_file_type(file_path)

    ocr_url = os.environ.get(
        "OCR_API_URL"
    )

    try:
        async with httpx.AsyncClient(timeout=300) as client:
            resp = await client.post(
                ocr_url,
                json={"image": b64_data, "file_type": file_type},
            )
            resp.raise_for_status()
            data = resp.json()
    except httpx.HTTPStatusError as e:
        raise ValueError(
            f"OCR API 返回错误: HTTP {e.response.status_code}"
        ) from e
    except httpx.RequestError as e:
        raise ValueError(
            f"OCR API 请求失败: {e}。请检查 OCR_API_URL 配置 ({ocr_url})"
        ) from e

    if isinstance(data, str):
        return data
    return data.get("result", str(data))


def main():
    if not os.environ.get("LLM_BASE_URL"):
        print(
            "错误: 环境变量 LLM_BASE_URL 未配置，请设置 Qwen VL API 地址",
            file=sys.stderr,
        )
        sys.exit(1)
    mcp.run()


if __name__ == "__main__":
    main()
