import asyncio
import unittest

from src.aiel_sdk.decorators import agent, router
from src.aiel_sdk.errors import ContractViolation
from src.aiel_sdk.exports import AgentExport, RouterExport, ToolExport
from src.aiel_sdk.registry import get_registry, reset_registry
from src.aiel_sdk.validate import validate_contracts
from src.aiel_sdk.adapters.langgraph.registry import to_langgraph_node, to_langgraph_router


class _Runtime:
    def __init__(self, context):
        self.context = context


class RouterRegistryTests(unittest.TestCase):
    def setUp(self):
        reset_registry()

    def test_router_registers_with_metadata_and_attaches_definition(self):
        @router("route_after_triage", description="Route triage", version="1.1", tags=("triage", "support"))
        def route_after_triage(ctx, state):
            _ = ctx
            return "research_agent" if state.get("route") == "research" else "response_agent"

        reg = get_registry()
        self.assertIn("route_after_triage", reg.routers)
        definition = reg.routers["route_after_triage"]
        self.assertEqual(definition.description, "Route triage")
        self.assertEqual(definition.version, "1.1")
        self.assertEqual(definition.tags, ("triage", "support"))
        self.assertIs(getattr(route_after_triage, "__aiel_router_definition__"), definition)

    def test_router_decorator_validates_signature(self):
        with self.assertRaises(ContractViolation) as ex:
            @router("route_bad")
            def route_bad(state):
                return "x"

        self.assertEqual(
            str(ex.exception),
            "router:route_bad must accept at least 2 args (ctx, state). Got: (state)",
        )

    def test_validate_contracts_error_messages_are_explicit(self):
        reg = get_registry()

        def bad_tool(ctx):
            return ctx

        def bad_agent(state):
            return state

        def bad_router(ctx):
            return "x"

        reg.tools["postgres.query"] = ToolExport(
            name="postgres.query",
            fn=bad_tool,
            module=bad_tool.__module__,
            qualname=bad_tool.__qualname__,
            is_async=False,
        )
        reg.agents["lookup_case"] = AgentExport(
            name="lookup_case",
            fn=bad_agent,
            module=bad_agent.__module__,
            qualname=bad_agent.__qualname__,
            is_async=False,
        )
        reg.routers["route_after_triage"] = RouterExport(
            name="route_after_triage",
            fn=bad_router,
            module=bad_router.__module__,
            qualname=bad_router.__qualname__,
            is_async=False,
        )

        with self.assertRaises(ContractViolation) as ex_tool:
            validate_contracts()
        self.assertEqual(
            str(ex_tool.exception),
            "tool:postgres.query must accept at least 2 args (ctx, input). Got: (ctx)",
        )

        reg.tools.clear()
        with self.assertRaises(ContractViolation) as ex_agent:
            validate_contracts()
        self.assertEqual(
            str(ex_agent.exception),
            "agent:lookup_case must accept at least 2 args (ctx, state/payload). Got: (state)",
        )

        reg.agents.clear()
        with self.assertRaises(ContractViolation) as ex_router:
            validate_contracts()
        self.assertEqual(
            str(ex_router.exception),
            "router:route_after_triage must accept at least 2 args (ctx, state). Got: (ctx)",
        )


class LangGraphRegistryAdapterTests(unittest.TestCase):
    def setUp(self):
        reset_registry()

    def test_to_langgraph_node_and_router(self):
        @agent("lookup_case")
        async def lookup_case(ctx, state):
            return {"case_exists": state["case_id"] == "INC-1", "tenant": ctx["tenant"]}

        @router("route_after_triage")
        def route_after_triage(ctx, state):
            _ = ctx
            return "research_agent" if state.get("route") == "research" else "response_agent"

        node = to_langgraph_node("lookup_case")
        route = to_langgraph_router("route_after_triage")
        runtime = _Runtime(context={"tenant": "t1"})

        node_out = asyncio.run(node({"case_id": "INC-1"}, runtime))
        route_out = asyncio.run(route({"route": "research"}, runtime))

        self.assertEqual(node_out["case_exists"], True)
        self.assertEqual(node_out["tenant"], "t1")
        self.assertEqual(route_out, "research_agent")

    def test_unknown_registry_entries_raise(self):
        with self.assertRaises(KeyError):
            _ = to_langgraph_node("missing_agent")
        with self.assertRaises(KeyError):
            _ = to_langgraph_router("missing_router")


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