"""Configuration 3: Hierarchical delegation.

One manager agent decomposes the task and delegates to two workers.
Comparable to McEntire's hierarchical condition (36% failure).
"""

from __future__ import annotations

import time

from cohort.benchmark.code_runner import extract_code, run_code_against_test
from cohort.benchmark.config import StrategyResult
from cohort.benchmark.strategies.base import Strategy

MANAGER_SYSTEM = """\
You are a team lead for a competitive programming team. You have two specialists:
- solver_a: Implements core algorithms and data structures
- solver_b: Handles edge cases, testing, and debugging

Your job is to:
1. Analyze the problem
2. Break it into subtasks
3. Assign each subtask to the right specialist
4. Integrate their work into a final solution

Format your delegation as:
TASK FOR solver_a: <description>
TASK FOR solver_b: <description>
"""

MANAGER_INTEGRATE_SYSTEM = """\
You are a team lead integrating your specialists' work into a final solution.
Combine solver_a's algorithm implementation with solver_b's edge case handling.
Output a single complete Python program (```python ... ```) that reads from stdin
and writes to stdout.
"""

SOLVER_A_SYSTEM = """\
You are a specialist competitive programmer focused on algorithm design and implementation.
Implement the assigned subtask in Python. Output a complete ```python ... ``` code block
that reads from stdin and writes to stdout.
"""

SOLVER_B_SYSTEM = """\
You are a specialist focused on edge cases, input validation, and debugging.
Review the problem constraints and write code or analysis addressing the assigned subtask.
If asked to handle edge cases, output a complete ```python ... ``` code block.
If asked to review, explain potential issues and fixes.
"""


class Hierarchical(Strategy):
    """Hierarchical: manager + 2 workers with delegation."""

    NAME = "hierarchical"

    def run(self) -> StrategyResult:
        t0 = time.time()
        failures: list[str] = []
        problem_text = self.task.to_prompt()

        # Phase 1: Manager decomposes
        plan = self._call_llm(
            agent_id="manager",
            system=MANAGER_SYSTEM,
            user=f"Decompose this problem and delegate subtasks:\n\n{problem_text}",
        )

        # Phase 2: Workers execute (sequentially -- they see each other's work)
        worker_a_response = self._call_llm(
            agent_id="solver_a",
            system=SOLVER_A_SYSTEM,
            user=(
                f"## Problem\n{problem_text}\n\n"
                f"## Your Assignment from Team Lead\n{plan}\n\n"
                f"Implement your assigned subtask."
            ),
        )

        worker_b_response = self._call_llm(
            agent_id="solver_b",
            system=SOLVER_B_SYSTEM,
            user=(
                f"## Problem\n{problem_text}\n\n"
                f"## Team Lead's Plan\n{plan}\n\n"
                f"## solver_a's Implementation\n{worker_a_response}\n\n"
                f"Handle your assigned subtask. Review solver_a's work for edge cases."
            ),
        )

        # Phase 3: Manager integrates
        integration = self._call_llm(
            agent_id="manager",
            system=MANAGER_INTEGRATE_SYSTEM,
            user=(
                f"## Problem\n{problem_text}\n\n"
                f"## solver_a's Work\n{worker_a_response}\n\n"
                f"## solver_b's Work\n{worker_b_response}\n\n"
                f"Integrate into a final solution."
            ),
        )

        code = extract_code(integration)

        # Phase 4: Optional fix -- if sample tests fail, delegate fix
        if code:
            errors = []
            for i, tc in enumerate(self.task.sample_cases):
                result = run_code_against_test(code, tc, self.task.time_limit_seconds)
                if not result.passed:
                    errors.append(
                        f"Test {i + 1}: input={tc.input!r}, "
                        f"expected={tc.expected_output!r}, "
                        f"got={result.actual!r}"
                        + (f" (error: {result.error})" if result.error else "")
                    )

            if errors:
                error_text = "\n".join(errors)
                fix_response = self._call_llm(
                    agent_id="solver_a",
                    system=SOLVER_A_SYSTEM,
                    user=(
                        f"## Problem\n{problem_text}\n\n"
                        f"## Current Solution\n```python\n{code}\n```\n\n"
                        f"## Test Failures\n{error_text}\n\n"
                        f"Fix the solution. Output the complete corrected code."
                    ),
                )
                fixed_code = extract_code(fix_response)
                if fixed_code:
                    code = fixed_code

        if not code:
            failures.append("no_code_produced")
            # Check if manager got stuck planning
            if not extract_code(worker_a_response) and not extract_code(worker_b_response):
                failures.append("planning_paralysis")

        elapsed = time.time() - t0
        return self._build_result(code, elapsed, failures)
