from __future__ import annotations

from ai_orchestrator.schemas.workflow_models import SubtaskRecord


class SupervisorAgent:
    def generate_task_plan(self, plan_steps: list[str]) -> list[str]:
        return plan_steps

    def create_subtasks(self, task_id: str, plan_steps: list[str]) -> list[SubtaskRecord]:
        subtasks: list[SubtaskRecord] = []
        for index, step in enumerate(plan_steps, start=1):
            subtasks.append(
                SubtaskRecord(
                    task_id=f"{task_id}_subtask_{index}",
                    parent_task=task_id,
                    description=step,
                    tool=self._tool_for_step(step),
                    expected_output="Completed step output",
                )
            )
        return subtasks

    def schedule_subtasks(self, subtasks: list[SubtaskRecord]) -> list[SubtaskRecord]:
        return subtasks

    def aggregate_results(self, *, plan_steps: list[str], failed_checks: list[str], issues: list[str]) -> dict:
        return {
            "plan_steps": plan_steps,
            "failed_checks": failed_checks,
            "issues": issues,
        }

    @staticmethod
    def _tool_for_step(step: str) -> str:
        normalized = step.lower()
        if "test" in normalized or "validate" in normalized:
            return "test_runner"
        if "inspect" in normalized or "repo" in normalized:
            return "repo_analyzer"
        if "patch" in normalized or "fix" in normalized:
            return "git_runner"
        return "shell_runner"
