import base64
import json
from pathlib import Path
from typing import Any

from app.providers.model_provider import ModelProvider


class CapabilityProbeService:
    def __init__(self, provider: ModelProvider):
        self.provider = provider

    async def probe(self) -> dict[str, Any]:
        result: dict[str, Any] = {}

        text = await self.provider.chat([{"role": "user", "content": "只回复 TEXT_OK"}])
        result["text_chat"] = "TEXT_OK" in text.content

        system = await self.provider.chat([
            {"role": "system", "content": "只回复 SYSTEM_OK"},
            {"role": "user", "content": "回复 USER_OK"},
        ])
        result["system_message"] = "SYSTEM_OK" in system.content

        multi = await self.provider.chat([
            {"role": "user", "content": "记住编号4837"},
            {"role": "assistant", "content": "已记住"},
            {"role": "user", "content": "编号是什么？只回复数字"},
        ])
        result["multi_turn"] = "4837" in multi.content

        result["json_output"] = await self._probe_json()
        result["tool_calling"] = await self._probe_tool()
        result["multimodal_image"] = await self._probe_image()
        return result

    async def _probe_json(self) -> str:
        reply = await self.provider.chat([
            {"role": "user", "content": '只输出 {"status":"ok"}'}
        ])
        try:
            value = json.loads(reply.content.strip())
            return "strict" if value.get("status") == "ok" else "none"
        except Exception:
            return "recoverable" if '"status"' in reply.content else "none"

    async def _probe_tool(self) -> str:
        tool = {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "查询天气",
                "parameters": {
                    "type": "object",
                    "properties": {"city": {"type": "string"}},
                    "required": ["city"],
                },
            },
        }
        reply = await self.provider.chat(
            [{"role": "user", "content": "请使用工具查询杭州天气"}],
            tools=[tool],
        )
        return "full" if reply.tool_calls else "none"

    async def _probe_image(self) -> bool:
        image_path = Path("data/capability_probe/blue_circle.png")
        if not image_path.exists():
            return False
        encoded = base64.b64encode(image_path.read_bytes()).decode("ascii")
        try:
            reply = await self.provider.chat([{
                "role": "user",
                "content": [
                    {"type": "text", "text": "图片是什么颜色和形状？"},
                    {"type": "image_url", "image_url": {
                        "url": f"data:image/png;base64,{encoded}"
                    }},
                ],
            }])
        except Exception:
            return False
        text = reply.content.lower()
        return ("蓝" in text or "blue" in text) and ("圆" in text or "circle" in text)