from __future__ import annotations

import unittest
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

from models import RuntimeModelResource
from services.vision import (
    VISION_ANALYSIS_SYSTEM_PROMPT,
    VISION_SYSTEM_HINT,
    generate_vision_focus_prompt,
    vision_analysis_system_prompt,
    vision_system_hint,
    vision_tool_description,
)
from services.workflow_engine import WorkflowEngine


class VisionFocusPromptTest(unittest.TestCase):
    def test_focus_extends_but_does_not_replace_safety_prompts(self):
        focus = "优先识别车辆数量和人员状态，只报告可确认的信息。"

        system_hint = vision_system_hint(focus)
        analysis_prompt = vision_analysis_system_prompt(focus)

        self.assertIn(VISION_SYSTEM_HINT, system_hint)
        self.assertIn(focus, system_hint)
        self.assertIn(VISION_ANALYSIS_SYSTEM_PROMPT, analysis_prompt)
        self.assertIn(focus, analysis_prompt)
        self.assertIn("优先回答当前用户问题", analysis_prompt)

    def test_empty_focus_preserves_previous_runtime_prompts(self):
        self.assertEqual(vision_system_hint(), VISION_SYSTEM_HINT)
        self.assertEqual(
            vision_analysis_system_prompt(),
            VISION_ANALYSIS_SYSTEM_PROMPT,
        )

    def test_tool_description_compacts_and_limits_focus(self):
        description = vision_tool_description("车辆\n  数量 " * 300)

        self.assertIn("当前业务重点", description)
        self.assertNotIn("\n", description)
        self.assertLess(len(description), 700)

    def test_workflow_resolves_global_and_node_focus(self):
        graph = {
            "specVersion": 3,
            "settings": {
                "visionEnabled": True,
                "visionFocusPrompt": "全局关注目标",
            },
            "nodes": [
                {
                    "id": "global-agent",
                    "type": "agent",
                    "data": {"name": "全局 Agent", "prompt": "处理任务"},
                },
                {
                    "id": "custom-agent",
                    "type": "agent",
                    "data": {
                        "name": "独立 Agent",
                        "prompt": "检查车辆",
                        "inheritGlobalConfig": False,
                        "visionEnabled": True,
                        "visionFocusPrompt": "节点关注目标",
                    },
                },
            ],
            "edges": [],
        }
        engine = WorkflowEngine(graph)

        self.assertEqual(
            engine.agent_stage_config("global-agent").vision_focus_prompt,
            "全局关注目标",
        )
        self.assertEqual(
            engine.agent_stage_config("custom-agent").vision_focus_prompt,
            "节点关注目标",
        )


class VisionFocusGenerationTest(unittest.IsolatedAsyncioTestCase):
    async def test_generates_editable_focus_with_selected_resource(self):
        resource = RuntimeModelResource(
            id="llm_1",
            name="生成模型",
            capability="LLM",
            interface_type="openai-llm",
            values={
                "apiUrl": "https://llm.example/v1",
                "modelId": "focus-model",
            },
            secrets={"apiKey": "secret"},
        )
        client = MagicMock()
        client.chat.completions.create = AsyncMock(
            return_value=SimpleNamespace(
                choices=[
                    SimpleNamespace(
                        message=SimpleNamespace(
                            content="优先识别车辆数量、损伤位置和人员状态。"
                        )
                    )
                ]
            )
        )
        client.close = AsyncMock()

        with patch("services.vision.AsyncOpenAI", return_value=client):
            result = await generate_vision_focus_prompt(
                resource,
                "你是交通事故处理助手，需要了解车辆和人员情况。",
            )

        self.assertEqual(result, "优先识别车辆数量、损伤位置和人员状态。")
        request = client.chat.completions.create.await_args.kwargs
        self.assertEqual(request["model"], "focus-model")
        self.assertIn("交通事故处理助手", request["messages"][1]["content"])
        client.close.assert_awaited_once()


if __name__ == "__main__":
    unittest.main()