"""Agent OS task scheduler runtime."""

from dataclasses import dataclass, field
from typing import Dict, List


@dataclass
class Task:
    task_id: str
    capability: str
    priority: int = 0
    status: str = "queued"
    assigned_agent: str = ""


class TaskScheduler:
    def __init__(self):
        self.tasks: Dict[str, Task] = {}
        self.agents: Dict[str, List[str]] = {}

    def register_agent(self, agent_id: str, capabilities: List[str]):
        self.agents[agent_id] = capabilities

    def submit(self, task: Task):
        self.tasks[task.task_id] = task

    def dispatch(self):
        for task in self.tasks.values():
            if task.status != "queued":
                continue
            for agent, caps in self.agents.items():
                if task.capability in caps:
                    task.assigned_agent = agent
                    task.status = "assigned"
                    break
        return list(self.tasks.values())