from typing import Literal, Dict, Any
from src.agents.code_agent import generate_code
from src.agents.vision_agent import generate_image
from src.agents.bot_agent import create_bot
from src.agents.reasoning_agent import plan_tasks


TaskType = Literal["code", "image", "bot", "plan"]


def run_task(task_type: TaskType, payload: Dict[str, Any]):
    if task_type == "code":
        return generate_code(payload.get("prompt", ""), payload.get("language", "python"))
    elif task_type == "image":
        return generate_image(payload.get("prompt", ""))
    elif task_type == "bot":
        return create_bot(payload.get("bot_type", "telegram"))
    elif task_type == "plan":
        return plan_tasks(payload.get("goal", ""))
    else:
        raise ValueError("Unsupported task type")