# echo_cog/se_state.py
"""
Minimal symbolic environment state for Echo Cog.

This is NOT the full agentic cell system; it's the thin slice needed
for 'Echo Cog MMLU' steering:
  - per-task accuracy tracking
  - list of last few Q/A for inclusion in prompts
"""

from __future__ import annotations

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


@dataclass
class QuestionRecord:
    task: str
    question: str
    choices: List[str]
    correct_answer: str  # "A"/"B"/"C"/"D"
    model_answer: Optional[str] = None
    correct: Optional[bool] = None


@dataclass
class TaskStats:
    asked: int = 0
    correct: int = 0

    @property
    def accuracy(self) -> float:
        return (self.correct / self.asked) if self.asked else 0.0


@dataclass
class SeState:
    """
    Minimal symbolic environment state for Echo Cog.

    Tracks:
      - per-task accuracy
      - list of last few Q/A for inclusion in prompts
    """

    history: List[QuestionRecord] = field(default_factory=list)
    task_stats: Dict[str, TaskStats] = field(default_factory=dict)
    max_history_per_task: int = 8

    def record(
        self,
        task: str,
        question: str,
        choices: List[str],
        correct_answer: str,
        model_answer: str,
    ) -> None:
        correct = (model_answer == correct_answer)
        rec = QuestionRecord(
            task=task,
            question=question,
            choices=choices,
            correct_answer=correct_answer,
            model_answer=model_answer,
            correct=correct,
        )
        self.history.append(rec)

        stats = self.task_stats.setdefault(task, TaskStats())
        stats.asked += 1
        if correct:
            stats.correct += 1

    def recent_for_task(self, task: str) -> List[QuestionRecord]:
        items = [r for r in self.history if r.task == task]
        return items[-self.max_history_per_task:]

    def summary_for_task(self, task: str) -> str:
        """
        Build a compact textual summary of prior performance for use in prompts.
        """
        stats = self.task_stats.get(task)
        if not stats or stats.asked == 0:
            return f"No prior questions seen for task '{task}'. Answer carefully."

        acc_pct = stats.accuracy * 100.0
        lines = [f"Task '{task}': {stats.correct}/{stats.asked} correct so far ({acc_pct:.1f}%)."]

        recent = self.recent_for_task(task)
        if not recent:
            return " ".join(lines)

        lines.append("Recent Q/A history (most recent last):")
        # Only include high-level question stems; avoid dumping full text walls
        for r in recent:
            yn = "+" if r.correct else "-"
            stem = r.question.strip().replace("\n", " ")
            if len(stem) > 120:
                stem = stem[:117] + "..."
            lines.append(
                f"  {yn} Q: {stem}  | correct={r.correct_answer}, model={r.model_answer or '?'}"
            )

        return "\n".join(lines)

    def get_recent_qa(self, task: str, k: int = 5) -> List[QuestionRecord]:
        """
        Return up to k most recent Q/A pairs for this task, oldest-first.

        This is used by EchoCogRuntime in 'rich' mode to include history context.
        """
        items = [r for r in self.history if r.task == task]
        return items[-k:] if k > 0 else []

    def stats_for_task(self, task: str) -> Dict[str, float]:
        """Return stats dict for a specific task (for SeController adapter)."""
        stats = self.task_stats.get(task)
        if stats is None:
            return {"accuracy": 0.5, "asked": 0, "correct": 0}
        return {
            "accuracy": stats.accuracy,
            "asked": stats.asked,
            "correct": stats.correct,
        }

    def overall_stats(self) -> Dict[str, float]:
        """Return overall accuracy across all tasks."""
        total_asked = sum(s.asked for s in self.task_stats.values())
        total_correct = sum(s.correct for s in self.task_stats.values())
        return {
            "total_asked": total_asked,
            "total_correct": total_correct,
            "accuracy": (total_correct / total_asked) if total_asked else 0.0,
        }
