"""Data classes for benchmark configuration, tasks, and results."""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class TestCase:
    """A single input/output test case."""

    input: str
    expected_output: str
    is_sample: bool = True  # visible to agents during solving


@dataclass
class CodingTask:
    """A competitive programming problem."""

    task_id: str
    title: str
    description: str
    input_format: str
    output_format: str
    constraints: str
    test_cases: list[TestCase] = field(default_factory=list)
    time_limit_seconds: float = 5.0
    difficulty: str = "medium"
    tags: list[str] = field(default_factory=list)

    @property
    def sample_cases(self) -> list[TestCase]:
        return [tc for tc in self.test_cases if tc.is_sample]

    @property
    def hidden_cases(self) -> list[TestCase]:
        return [tc for tc in self.test_cases if not tc.is_sample]

    def to_prompt(self) -> str:
        """Format the problem as a prompt for agents."""
        parts = [
            f"# {self.title}\n",
            self.description,
            f"\n## Input Format\n{self.input_format}",
            f"\n## Output Format\n{self.output_format}",
            f"\n## Constraints\n{self.constraints}",
            "\n## Sample Test Cases",
        ]
        for i, tc in enumerate(self.sample_cases, 1):
            parts.append(f"\n### Sample {i}")
            parts.append(f"**Input:**\n```\n{tc.input}\n```")
            parts.append(f"**Output:**\n```\n{tc.expected_output}\n```")
        return "\n".join(parts)


@dataclass
class TokenUsage:
    """Accumulated token usage across a strategy run."""

    tokens_in: int = 0
    tokens_out: int = 0
    api_calls: int = 0

    @property
    def total_tokens(self) -> int:
        return self.tokens_in + self.tokens_out

    def record(self, tokens_in: int, tokens_out: int) -> None:
        self.tokens_in += tokens_in
        self.tokens_out += tokens_out
        self.api_calls += 1


@dataclass
class TurnLog:
    """Record of a single agent turn."""

    turn_number: int
    agent_id: str
    content: str
    tokens_in: int = 0
    tokens_out: int = 0
    relevance_score: float | None = None
    phase: str = ""


@dataclass
class StrategyResult:
    """Output from running one strategy on one task."""

    strategy_name: str
    final_code: str
    turns: list[TurnLog] = field(default_factory=list)
    token_usage: TokenUsage = field(default_factory=TokenUsage)
    elapsed_seconds: float = 0.0
    failure_modes: list[str] = field(default_factory=list)
    conversation_log: list[dict] = field(default_factory=list)


@dataclass
class TaskMetrics:
    """Computed metrics for one strategy on one task."""

    strategy_name: str
    task_id: str

    # Correctness
    sample_passed: int = 0
    sample_total: int = 0
    hidden_passed: int = 0
    hidden_total: int = 0

    # Cost
    total_tokens: int = 0
    api_calls: int = 0

    # Efficiency
    turn_count: int = 0
    elapsed_seconds: float = 0.0

    # Failure modes
    loop_count: int = 0
    duplication_count: int = 0
    planning_paralysis: bool = False
    off_topic_turns: int = 0

    @property
    def correctness(self) -> float:
        total = self.sample_total + self.hidden_total
        if total == 0:
            return 0.0
        return (self.sample_passed + self.hidden_passed) / total

    def to_dict(self) -> dict:
        return {
            "strategy": self.strategy_name,
            "task_id": self.task_id,
            "correctness": round(self.correctness, 4),
            "sample_passed": self.sample_passed,
            "sample_total": self.sample_total,
            "hidden_passed": self.hidden_passed,
            "hidden_total": self.hidden_total,
            "total_tokens": self.total_tokens,
            "api_calls": self.api_calls,
            "turn_count": self.turn_count,
            "elapsed_seconds": round(self.elapsed_seconds, 2),
            "loop_count": self.loop_count,
            "duplication_count": self.duplication_count,
            "planning_paralysis": self.planning_paralysis,
            "off_topic_turns": self.off_topic_turns,
        }
