from dataiku.llm.agent_tools import BaseAgentTool
from shortcut_api import pick, shortcut_get


class ListWorkflows(BaseAgentTool):
    def get_descriptor(self, tool):
        return {
            "description": "List Shortcut workflows and their state IDs, names and types.",
            "inputSchema": {"type": "object", "properties": {}, "additionalProperties": False},
        }

    def invoke(self, input, trace):
        workflows = []
        for workflow in shortcut_get("workflows"):
            workflows.append({
                **pick(workflow, ("id", "name")),
                "states": [pick(state, ("id", "name", "type")) for state in workflow["states"]],
            })
        return {"output": workflows}