"""Configuration 1: Single agent baseline.

One agent, one prompt, optional fix iterations. The simplest approach --
directly comparable to McEntire's single-agent condition (28/28 success).
"""

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

SYSTEM_PROMPT = """\
You are an expert competitive programmer. Solve the given problem in Python.

Rules:
- Read input from stdin, write output to stdout
- Your solution must handle all edge cases within the given constraints
- Output ONLY a single Python code block (```python ... ```)
- No explanations needed -- just the code
"""

FIX_PROMPT = """\
Your previous solution failed on some test cases.

{errors}

Fix the code. Output ONLY the corrected Python code block.
"""


class SingleAgent(Strategy):
    """Single agent: one shot + up to 2 fix attempts."""

    NAME = "single_agent"
    MAX_FIX_ATTEMPTS = 2

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

        # Initial solve
        response = self._call_llm(
            agent_id="solver",
            system=SYSTEM_PROMPT,
            user=self.task.to_prompt(),
        )
        code = extract_code(response)

        # Fix loop: check sample cases, fix if needed
        for attempt in range(self.MAX_FIX_ATTEMPTS):
            if not code:
                failures.append("no_code_extracted")
                break

            # Check sample cases
            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 not errors:
                break  # All sample cases pass

            # Ask agent to fix
            error_text = "\n".join(errors)
            fix_response = self._call_llm(
                agent_id="solver",
                system=SYSTEM_PROMPT,
                user=FIX_PROMPT.format(errors=error_text),
            )
            code = extract_code(fix_response)

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