from __future__ import annotations

import unittest

from services.runtime_variables import DynamicVariableStore
from services.workflow.models import LLMRouteResult, RouteStatus
from services.workflow.routing import WorkflowEdgeEvaluator
from services.workflow_engine import WorkflowEngine


def routing_graph() -> dict:
    return {
        "specVersion": 3,
        "settings": {},
        "nodes": [
            {"id": "start", "type": "start", "data": {}},
            {"id": "agent", "type": "agent", "data": {"name": "Agent"}},
            {"id": "matched", "type": "end", "data": {}},
            {"id": "fallback", "type": "end", "data": {}},
        ],
        "edges": [
            {
                "id": "condition",
                "source": "agent",
                "target": "matched",
                "data": {
                    "mode": "llm",
                    "priority": 10,
                    "condition": "用户要求结束",
                },
            },
            {
                "id": "default",
                "source": "agent",
                "target": "fallback",
                "data": {"mode": "always", "priority": 10},
            },
        ],
    }


class WorkflowEdgeEvaluatorTest(unittest.IsolatedAsyncioTestCase):
    async def test_router_error_does_not_take_default_path(self):
        class ErrorRouter:
            async def select_edge(self, **_kwargs):
                return LLMRouteResult(
                    status=RouteStatus.ERROR,
                    error="provider unavailable",
                )

        evaluator = WorkflowEdgeEvaluator(
            WorkflowEngine(routing_graph()),
            DynamicVariableStore({}),
            lambda _node_id: ErrorRouter(),
        )

        decision = await evaluator.evaluate("agent")

        self.assertEqual(decision.status, RouteStatus.ERROR)
        self.assertIsNone(decision.edge)

    async def test_explicit_no_match_takes_default_path(self):
        class NoMatchRouter:
            async def select_edge(self, **_kwargs):
                return LLMRouteResult(status=RouteStatus.NO_MATCH)

        evaluator = WorkflowEdgeEvaluator(
            WorkflowEngine(routing_graph()),
            DynamicVariableStore({}),
            lambda _node_id: NoMatchRouter(),
        )

        decision = await evaluator.evaluate("agent")

        self.assertEqual(decision.status, RouteStatus.MATCHED)
        self.assertEqual(decision.edge["id"], "default")


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