import ray
from pydantic_ai import Agent
from pretor.core.workflow_manager.workflow import WorkStep

@ray.remote
class ControlNode:
    def __init__(self, agent: Agent):
        self.agent = agent

    async def execute_step(self, step: WorkStep) -> WorkStep:
        if step.action == "dispatch_model":
            # The WorkStep schema from workflow manager may pass target info differently, assuming `input` here or simple `desc`
            result = await self.dispatch_model({}, f"Execute task: {step.desc}")
            step.output = result
        elif step.action == "dispatch_tool":
            # Simulating parsing of tool and args from `desc` or `input`
            result = await self.dispatch_tool("simulated_tool", {"desc": step.desc})
            step.output = str(result)
        else:
            result = await self.agent.run(f"Execute generic step: {step.model_dump()}")
            step.output = result.data

        step.status = "completed"
        return step

    async def dispatch_model(self, model_info: dict, prompt: str) -> str:
        # In a real system, we'd select a smaller/specific model based on model_info
        result = await self.agent.run(prompt)
        return result.data

    async def dispatch_tool(self, tool_name: str, tool_args: dict) -> dict:
        # Simulated tool dispatch
        return {"tool": tool_name, "status": "executed", "args": tool_args}