"""Read the canonical session text and recover what the grader will ask about.

The arena renders every session the same way:

    [user]
    [assistant]
    [tool_use kind=edit name=Write path=src/pkg/a.py]
    [tool_result ok lines=12 kept=12]

The grader's facts (files, exception types, symbols) and its continuation
probes (last edited file, last test result, last exception, last traceback
function, created file) are all derived from that history. This module walks
the rendered text the same way the arena's observer walks the raw session, so
the harness can make sure those exact answers are in the summary.

Pure standard library: the sandbox allows no extra packages for this.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field

MAX_FILES = 12
MAX_ERRORS = 5
MAX_SYMBOLS = 8

HEADER = re.compile(
    r"^\[(user|assistant|tool_use|tool_result)((?:\s[^\]\n]*)?)\]\s*$", re.MULTILINE
)
ATTR = re.compile(r"(\w+)=(\S+)")
FRAME = re.compile(r'File "([^"]+)", line \d+, in ([A-Za-z_]\w*)')
PYTEST_FRAME = re.compile(r"^([\w./-]+\.py):\d+: in ([A-Za-z_]\w*)", re.MULTILINE)
ERROR_LINE = re.compile(
    r"^\s*(?:E\d*\s+)?([A-Za-z_][\w.]*(?:Error|Exception|Panic|Fault))\b(?::\s*(.*))?$",
    re.MULTILINE,
)
DEF = re.compile(
    r"^\s*(?:async\s+)?def\s+([A-Za-z_]\w*)\s*\(|^\s*class\s+([A-Za-z_]\w*)", re.MULTILINE
)
PYTEST_TAIL = re.compile(
    r"^([\w./-]+\.py):\d+: ([A-Za-z_]\w*(?:Error|Exception))\s*$", re.MULTILINE
)
PASSED = re.compile(r"\b(\d+) passed\b")
FAILED = re.compile(r"\b(\d+) (?:failed|error)s?\b")
MISSING = ("No such file or directory", "does not exist", "not found", "cannot access")
STATUS = {"ok": True, "error": False, "done": None}


@dataclass
class Block:
    kind: str
    status: str
    attrs: dict[str, str]
    body: str


@dataclass
class Observed:
    opened: list[str] = field(default_factory=list)
    created: list[str] = field(default_factory=list)
    edited: list[str] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)
    frames: list[str] = field(default_factory=list)
    symbols: list[str] = field(default_factory=list)
    tests: list[bool] = field(default_factory=list)
    users: list[str] = field(default_factory=list)
    last_kind: str = ""
    last_path: str = ""


def _add(found: list[str], value: str) -> None:
    if value and value not in found:
        found.append(value)


def _in_repo(path: str) -> bool:
    clean = path.replace("\\", "/")
    return not ("site-packages" in clean or "/lib/python" in clean or clean.startswith("/usr/"))


def _test_passed(text: str, ok: bool | None) -> bool:
    if FAILED.search(text) and not re.search(r"\b0 (?:failed|errors?)\b", text):
        return False
    if PASSED.search(text):
        return True
    if ok is not None:
        return ok
    lowered = text.lower()
    return " failed" not in lowered and "error" not in lowered


def blocks(user: str) -> list[Block]:
    text = user.replace("\r\n", "\n")
    text = re.sub(r"^\s*<session>\s*", "", text)
    text = re.sub(r"\s*</session>\s*$", "", text)
    heads = list(HEADER.finditer(text))
    found: list[Block] = []
    for index, head in enumerate(heads):
        end = heads[index + 1].start() if index + 1 < len(heads) else len(text)
        rest = head.group(2) or ""
        words = rest.split()
        status = words[0] if words and "=" not in words[0] else ""
        found.append(
            Block(head.group(1), status, dict(ATTR.findall(rest)), text[head.end() : end].strip("\n"))
        )
    return found


def observe(user: str) -> Observed:
    """Port of the arena's observer, working on the rendered session text."""
    seen = Observed()
    pending: list[tuple[str, str, str]] = []
    known: set[str] = set()
    previous = ""
    for block in blocks(user):
        if block.kind in ("assistant", "tool_use"):
            if previous not in ("assistant", "tool_use"):
                pending = []
            if block.kind == "tool_use":
                kind = block.attrs.get("kind", "")
                path = block.attrs.get("path", "")
                pending.append((kind, path, block.attrs.get("name", "")))
                seen.last_kind, seen.last_path = kind, path
                if kind == "edit":
                    for name in DEF.findall(block.body):
                        _add(seen.symbols, name[0] or name[1])
            previous = block.kind
            continue
        previous = block.kind
        text = block.body
        if block.kind == "user":
            seen.users.append(text)
        if block.kind == "tool_result" and pending:
            kind, path, name = pending.pop(0)
            ok = STATUS.get(block.status)
            lowered = text.lower()
            if path and _in_repo(path):
                if kind == "read" and not any(m in text for m in MISSING):
                    _add(seen.opened, path)
                elif kind == "edit" and ok is not False:
                    creates = name in ("Write", "create") or "created" in lowered
                    if path not in known and creates:
                        _add(seen.created, path)
                    else:
                        _add(seen.edited, path)
                known.add(path)
            if kind == "test":
                seen.tests.append(_test_passed(text, ok))
        for path, function in [*FRAME.findall(text), *PYTEST_FRAME.findall(text)]:
            if _in_repo(path) and function != "<module>":
                seen.frames.append(function)
                _add(seen.symbols, function)
        for name, _message in ERROR_LINE.findall(text):
            seen.errors.append(name)
        for _path, name in PYTEST_TAIL.findall(text):
            seen.errors.append(name)
    return seen


def recent_unique(values: list[str], limit: int) -> list[str]:
    ordered: list[str] = []
    for value in values:
        if value in ordered:
            ordered.remove(value)
        ordered.append(value)
    return ordered[-limit:]


def fact_files(seen: Observed) -> list[str]:
    files: list[str] = []
    for path in [*seen.opened, *seen.created, *seen.edited]:
        _add(files, path)
    keep = [p for p in files if p in seen.edited or p in seen.created]
    keep += [p for p in files if p not in keep]
    return sorted(keep[:MAX_FILES], key=files.index)


def facts(seen: Observed) -> dict[str, list[str]]:
    return {
        "files": fact_files(seen),
        "errors": recent_unique(seen.errors, MAX_ERRORS),
        "symbols": recent_unique(seen.symbols, MAX_SYMBOLS),
    }