#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = []
# ///
"""reflect_distill.py — mechanical distillation of session logs for the /reflect skill.

Turns `sessions/*.jsonl` into a compact, readable timeline the analysing LLM can scan.
Tool results are ~89% of the corpus by volume and carry almost no diagnostic value, so
they collapse to `name(args) -> ok|ERROR, size`; user and assistant prose is kept whole
because that is where intent shows.

This script makes **no quality judgements** — no error detectors, no ranking, no
"suspicious" flags. Deciding what is a mistake is the LLM's job; mechanical reduction
is this script's job. That split is deliberate (see plans/reflect-skill.md).
"""

import argparse
import base64
import binascii
import itertools
import json
import re
import sys
from collections.abc import Iterator
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

# workspace/skills/reflect/scripts/reflect_distill.py -> parents[3] = workspace root.
WORKSPACE = Path(__file__).resolve().parents[3]

# Session-key prefixes that are machinery or throwaway tests, not conversations worth
# diagnosing. `reflect` is here so the skill never analyses its own runs.
NOISE_PREFIXES = (
    "reflect",
    "compact-memory-auto",
    "detach",
    "dream",
    "cron",
    "cli",
    "wiki-compile",
    "wiki-capture-test",
    "note-compile",
    "test",
)

SESSION_MARKER = "━━━ SESSION"
MIN_MESSAGES = 5
REASONING_CHARS = 200
ARG_VALUE_CHARS = 90

# The context window is not what bounds a batch — latency is. Measured on real digests this
# corpus runs about 1.2 characters per token, so 500k chars was ~185k tokens per request, whose
# prefill on glm-5.3:cloud kept overrunning the provider timeout and each retry paid for it
# again (2026-09-02 run: 9 timeouts, two turns lost outright). 200k chars ≈ 70k tokens answers
# in one pass; smaller batches also get a more careful read and, since reflect_auto.py files
# findings after every batch, cost nothing but more commit points. It imports this rather than
# keeping its own copy, so a debugging run batches exactly like the real one.
DEFAULT_BUDGET_CHARS = 200_000

_WORKSPACE_PATH_RE = re.compile(r"/home/[^/]+/\.nanobot/workspace/")
_WHITESPACE_RE = re.compile(r"\s+")


@dataclass(frozen=True)
class SessionDigest:
    """One distilled session: identity, size, and the rendered timeline.

    `started` is the raw ISO timestamp exactly as the log holds it, because it doubles as
    the cursor `collect_sessions` compares against. Formatting happens only for display.
    """

    name: str
    started: str
    message_count: int
    text: str

    def __len__(self) -> int:
        return len(self.text)


def _decode_session_name(stem: str) -> str:
    """Return the readable session key.

    Older sessions are stored under a base64 filename (`ZHJlYW06...` = `dream:...`);
    decoding them is what lets prefix filtering catch that generation too.
    """
    if "_" in stem or "-" in stem:
        return stem
    padded = stem + "=" * (-len(stem) % 4)
    try:
        decoded = base64.urlsafe_b64decode(padded).decode("utf-8")
    except (binascii.Error, UnicodeDecodeError, ValueError):
        return stem
    return decoded if decoded.isprintable() else stem


def is_noise(session_name: str) -> bool:
    """True for machinery sessions that carry no diagnostic value.

    Prefix match, not exact: variants like `wiki-capture-test2` or `test-note-search-zzz`
    are the same throwaway machinery as their base name.
    """
    key = _decode_session_name(session_name)
    return key.startswith(NOISE_PREFIXES)


def _iter_records(path: Path) -> Iterator[dict]:
    """Yield the JSON records of a session log, skipping blank and malformed lines."""
    with path.open(encoding="utf-8", errors="replace") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError:
                continue


def _read_records(path: Path) -> list[dict]:
    return list(_iter_records(path))


def _format_stamp(raw: str) -> str:
    """Display form of the raw ISO timestamp: `2026-07-11 14:02`."""
    return raw[:16].replace("T", " ") if raw else "unknown"


def _shorten(text: str, limit: int) -> str:
    """Collapse whitespace and cut to `limit`, marking the cut with an ellipsis."""
    flat = _WHITESPACE_RE.sub(" ", text).strip()
    return flat if len(flat) <= limit else flat[:limit] + "…"


def _shorten_middle(text: str, limit: int) -> str:
    """Shorten from the middle, keeping both ends.

    Argument values are usually URLs and paths whose distinguishing part sits at the
    *end*; cutting from the front would render two different fetches identical and
    hide the difference between a genuine retry and legitimate sequential work.
    """
    flat = _WHITESPACE_RE.sub(" ", text).strip()
    if len(flat) <= limit:
        return flat
    head = (limit * 2) // 3
    tail = limit - head
    return f"{flat[:head]}…{flat[-tail:]}"


def _format_size(char_count: int) -> str:
    if char_count < 1000:
        return f"{char_count} B"
    return f"{char_count / 1000:.1f} kB"


def _format_arguments(raw: str) -> str:
    """Render tool arguments compactly, stripping the workspace path prefix.

    Falls back to the raw string when arguments are not JSON — some providers emit
    partial or malformed argument blobs, and that is itself worth seeing.
    """
    try:
        parsed = json.loads(raw) if raw else {}
    except json.JSONDecodeError:
        return _shorten_middle(raw, ARG_VALUE_CHARS)
    if not isinstance(parsed, dict):
        return _shorten_middle(str(parsed), ARG_VALUE_CHARS)
    parts = []
    for key, value in parsed.items():
        text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
        text = _WORKSPACE_PATH_RE.sub("", text)
        parts.append(f"{key}={_shorten_middle(text, ARG_VALUE_CHARS)}")
    return ", ".join(parts)


def _tool_outcome(content: str) -> str:
    """Classify a tool result as ok or ERROR, keeping the error's first line.

    Detection is textual because tool results carry no status field; the leading
    `Error:` / `Traceback` / `{"error"` shapes are what the runtime actually emits.
    """
    head = content.lstrip()[:80]
    if re.match(r'(?i)^(error|exception|traceback|failed|\{"error)', head):
        return f"ERROR {_shorten(content, ARG_VALUE_CHARS)}"
    return "ok"


def distill_session(path: Path) -> SessionDigest | None:
    """Render one session file as a timeline, or None when it should be skipped.

    Skipped: machinery sessions (see NOISE_PREFIXES) and sessions shorter than
    MIN_MESSAGES, which are too short to show a behavioural pattern.
    """
    if is_noise(path.stem):
        return None

    session_key = _decode_session_name(path.stem)
    records = _read_records(path)
    started = ""
    lines: list[str] = []
    message_count = 0
    pending_calls: dict[str, str] = {}

    for record in records:
        if record.get("_type") == "metadata":
            started = str(record.get("created_at") or "")
            continue
        # Slash commands are handled by the runtime, not by the agent's reasoning.
        if record.get("_command"):
            continue

        role = record.get("role")
        content = record.get("content") or ""
        if not isinstance(content, str):
            content = json.dumps(content, ensure_ascii=False)
        if not started:
            started = str(record.get("timestamp") or "")

        if role == "user":
            message_count += 1
            lines.append(f"u: {content.strip()}")
        elif role == "assistant":
            message_count += 1
            reasoning = record.get("reasoning_content") or ""
            if reasoning:
                lines.append(f"   ~ {_shorten(reasoning, REASONING_CHARS)}")
            if content.strip():
                lines.append(f"a: {content.strip()}")
            for call in record.get("tool_calls") or []:
                function = call.get("function") or {}
                name = function.get("name") or "?"
                pending_calls[call.get("id") or ""] = name
                lines.append(f"a: → {name}({_format_arguments(function.get('arguments') or '')})")
        elif role == "tool":
            message_count += 1
            name = record.get("name") or pending_calls.get(record.get("tool_call_id") or "", "?")
            lines.append(f"   ← {name}: {_tool_outcome(content)}, {_format_size(len(content))}")

    if message_count < MIN_MESSAGES:
        return None

    # Box-drawing marker, not a markdown heading: assistant prose is full of `###`,
    # so a heading would not read as a session boundary.
    header = f"{SESSION_MARKER} {session_key} | {_format_stamp(started)} | {message_count} messages"
    return SessionDigest(
        name=session_key,
        started=started,
        message_count=message_count,
        text="\n".join([header, *lines]),
    )


def _session_start(path: Path) -> str:
    """First timestamp inside the file — mtime is unreliable after bulk file moves.

    Reads only the opening records: this runs for every session on every run, and parsing
    whole files here would mean reading the entire corpus twice.
    """
    for record in itertools.islice(_iter_records(path), 5):
        stamp = record.get("created_at") or record.get("timestamp")
        if stamp:
            return str(stamp)
    return ""


def collect_sessions(sessions_dir: Path, since: str = "") -> list[Path]:
    """Session files newer than `since` (ISO timestamp), oldest first."""
    paths = []
    for path in sorted(sessions_dir.glob("*.jsonl")):
        if is_noise(path.stem):
            continue
        start = _session_start(path)
        if since and start and start <= since:
            continue
        paths.append((start, path))
    return [path for _, path in sorted(paths)]


def iter_batches(paths: list[Path], budget_chars: int) -> Iterator[list[SessionDigest]]:
    """Group distilled sessions into batches that fit the per-turn character budget.

    A single session larger than the budget still gets its own batch — truncating it
    would hide exactly the long runaway loops worth finding.
    """
    batch: list[SessionDigest] = []
    size = 0
    for path in paths:
        digest = distill_session(path)
        if digest is None:
            continue
        if batch and size + len(digest) > budget_chars:
            yield batch
            batch, size = [], 0
        batch.append(digest)
        size += len(digest)
    if batch:
        yield batch


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--sessions-dir", type=Path, default=WORKSPACE / "sessions")
    parser.add_argument("--since", default="", help="ISO timestamp; only sessions started after it")
    parser.add_argument("--budget-chars", type=int, default=DEFAULT_BUDGET_CHARS)
    parser.add_argument("--out", type=Path, help="write batches to OUT.NNN.md instead of stdout")
    parser.add_argument("--stats", action="store_true", help="print a size summary to stderr")
    args = parser.parse_args(argv)

    if not args.sessions_dir.is_dir():
        print(f"sessions dir not found: {args.sessions_dir}", file=sys.stderr)
        return 1

    paths = collect_sessions(args.sessions_dir, args.since)
    batches = list(iter_batches(paths, args.budget_chars))

    for index, batch in enumerate(batches):
        body = "\n\n".join(digest.text for digest in batch)
        if args.out:
            target = args.out.with_suffix(f".{index:03d}.md")
            target.write_text(body, encoding="utf-8")
        else:
            print(body)

    if args.stats:
        sessions = sum(len(batch) for batch in batches)
        chars = sum(len(digest) for batch in batches for digest in batch)
        print(
            f"{datetime.now():%Y-%m-%d %H:%M}  candidates {len(paths)}, "
            f"distilled {sessions} sessions, {len(batches)} batches, {chars / 1000:.0f} kB",
            file=sys.stderr,
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())