import os
import unittest
from mcpuniverse.mcp.manager import MCPManager
from mcpuniverse.workflows.builder import WorkflowBuilder
from mcpuniverse.common.context import Context


class TestWorkflowBuilder(unittest.TestCase):

    def test_get_constructor_signature(self):
        builder = WorkflowBuilder(mcp_manager=MCPManager(), config={})
        params = builder._get_constructor_signature(builder._workflow_classes["router"])
        self.assertTrue(params["llm"]["is_llm_or_executor"])
        self.assertTrue(params["agents"]["is_llm_or_executor"])
        self.assertFalse(params["top_k"]["is_llm_or_executor"])

    def test_build_dependency_graph(self):
        folder = os.path.dirname(os.path.realpath(__file__))
        config_file = os.path.join(folder, "../data/config/yaml_config.yaml")
        workflow = WorkflowBuilder(mcp_manager=MCPManager(), config=config_file)
        graph = workflow._build_dependency_graph()
        self.assertListEqual(graph["agent-1"], ["llm-1"])
        self.assertListEqual(graph["agent-2"], ["llm-1"])
        self.assertListEqual(graph["agent-3"], ["llm-1"])
        self.assertListEqual(sorted(graph["workflow-1"]), ["agent-1", "agent-2", "agent-3"])

    def test_build_workflow(self):
        folder = os.path.dirname(os.path.realpath(__file__))
        config_file = os.path.join(folder, "../data/config/yaml_config.yaml")
        workflow = WorkflowBuilder(mcp_manager=MCPManager(), config=config_file)
        workflow.build()
        description = workflow.get_component("main").get_description()
        self.assertTrue("Main agent" in description)

    def test_dump_config(self):
        folder = os.path.dirname(os.path.realpath(__file__))
        config_file = os.path.join(folder, "../data/config/yaml_config.yaml")
        workflow = WorkflowBuilder(mcp_manager=MCPManager(), config=config_file)
        workflow.build(project_id="test")
        configs = workflow.dump_config()
        name2id = {}
        for c in configs:
            if c["kind"] == "agent":
                name2id[c["spec"]["name"]] = f'test:agent:{c["spec"]["name"]}'
            elif c["kind"] == "llm":
                name2id[c["spec"]["name"]] = f'test:llm:{c["spec"]["name"]}'
            elif c["kind"] == "workflow":
                name2id[c["spec"]["name"]] = f'test:workflow:{c["spec"]["name"]}'

        workflow = WorkflowBuilder(mcp_manager=MCPManager(), config=configs)
        workflow.build(project_id="test")
        for name in name2id.keys():
            self.assertEqual(workflow.get_component(name).id, name2id[name])

    def test_list_undefined_env_vars(self):
        folder = os.path.dirname(os.path.realpath(__file__))
        config_file = os.path.join(folder, "../data/config/yaml_config.yaml")
        workflow = WorkflowBuilder(mcp_manager=MCPManager(), config=config_file)
        workflow.build(project_id="test")

        os.environ["OPENAI_API_KEY"] = ""
        r = workflow.list_undefined_env_vars()
        self.assertListEqual(r, ["OPENAI_API_KEY"])
        context = Context(env={"OPENAI_API_KEY": "xxx"})
        workflow.set_context(context)
        r = workflow.list_undefined_env_vars()
        self.assertListEqual(r, [])

    def test_claude_code_config(self):
        folder = os.path.dirname(os.path.realpath(__file__))
        config_file = os.path.join(folder, "../data/config/claude_code.yaml")
        workflow = WorkflowBuilder(mcp_manager=MCPManager(), config=config_file)
        workflow.build(project_id="test")
        agent = workflow.get_component(name="agent-1")
        self.assertEqual(type(agent).__name__, "ClaudeCodeAgent")


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