"""Stream-based repair of bloated pi_session*.json files.

Two historical sources of session bloat:

1. ``message_update`` / ``assistantMessageEvent.type=toolcall_delta`` events,
   which are cumulative argument snapshots for in-flight tool calls.
2. Redundant cumulative snapshots embedded on *every* ``message_update`` under
   ``message`` and ``assistantMessageEvent.partial``.

Neither is required by abreka's session viewer, which renders from
``thinking_delta`` / ``text_delta`` plus message/tool boundaries.

This module provides a streaming repair primitive that drops/compacts these
payloads in place without ever fully loading a file during the scanning pass.

Safety:

- Scanning is read-only and streams line by line.
- Repair writes to ``<path>.repair.tmp``, validates, then ``os.replace``s
  atomically onto the original. The original file is never opened for writing.
- Validation rejects any temp file whose reloaded event count does not match
  ``plan.total_events - plan.dropped_events``.
- Files without any repairable content are a no-op (mtime untouched).
"""

from __future__ import annotations

import dataclasses
import json
import os
from pathlib import Path
from typing import Iterator

from abreka.pi_client import compact_event_for_session, is_dropped_event



@dataclasses.dataclass
class RepairPlan:
    path: Path
    original_size: int
    total_events: int
    dropped_events: int
    compacted_events: int
    estimated_saved_bytes: int
    kept_events: list[dict] | None
    ok: bool
    error: str | None = None

    @property
    def needs_repair(self) -> bool:
        return self.ok and (self.dropped_events > 0 or self.compacted_events > 0)


class SessionRepairError(Exception):
    """Raised when a repair operation cannot proceed safely."""


# ---------- streaming parser --------------------------------------------------


def _iter_events(path: Path) -> Iterator[dict]:
    """Yield events from a pi session file written with json.dump(indent=2).

    Expects the exact format produced by ``PiRpcClient.save_session``:

        [
          { ...event... },
          { ...event... }
        ]

    Peak memory is bounded by the single largest event.
    """
    with path.open("rb") as f:
        # Validate the first non-whitespace byte is '['.
        first = f.read(1)
        while first and first in b" \t\r\n":
            first = f.read(1)
        if first != b"[":
            raise SessionRepairError(
                f"{path}: does not begin with '[' — not a pi-rpc session array"
            )

        # Fast path: empty list `[]` (what json.dump(indent=2) produces for []).
        rest = f.read().lstrip(b" \t\r\n")
        if rest.startswith(b"]"):
            tail = rest[1:].lstrip(b" \t\r\n")
            if tail:
                raise SessionRepairError(
                    f"{path}: trailing bytes after ']': {tail[:40]!r}"
                )
            return

        # Rewind and parse by lines. json.dump(indent=2) on a non-empty list
        # produces:
        #   [
        #     { ...event... },
        #     { ...event... }
        #   ]
        f.seek(0)
        saw_open_bracket = False
        buf: bytearray | None = None

        for raw in f:
            if buf is None:
                stripped = raw.strip()
                if not saw_open_bracket:
                    if stripped == b"[":
                        saw_open_bracket = True
                        continue
                    if stripped.startswith(b"["):
                        raise SessionRepairError(
                            f"{path}: unexpected header {stripped!r}"
                        )
                    continue
                if stripped == b"]":
                    return
                if stripped == b"{":
                    buf = bytearray(b"{")
                    continue
                if stripped == b"":
                    continue
                raise SessionRepairError(
                    f"{path}: unexpected line outside event: {stripped[:80]!r}"
                )
            else:
                stripped = raw.rstrip(b"\r\n")
                if stripped in (b"  }", b"  },"):
                    buf.extend(b"}")
                    try:
                        yield json.loads(bytes(buf))
                    except json.JSONDecodeError as e:
                        raise SessionRepairError(
                            f"{path}: event JSON parse failed: {e}"
                        ) from e
                    buf = None
                else:
                    buf.extend(raw)

        if buf is not None:
            raise SessionRepairError(f"{path}: unterminated event at EOF")


def scan_session_file(path: Path, *, counts_only: bool = False) -> RepairPlan:
    """Stream-parse *path* and return a RepairPlan.

    With ``counts_only=True``, transformed events are counted but not retained
    in the returned plan — use this for dry-run summaries over large files.
    """
    try:
        st = path.stat()
    except OSError as e:
        return RepairPlan(
            path=path,
            original_size=0,
            total_events=0,
            dropped_events=0,
            compacted_events=0,
            estimated_saved_bytes=0,
            kept_events=None,
            ok=False,
            error=f"stat failed: {e}",
        )

    if not path.is_file():
        return RepairPlan(
            path=path,
            original_size=st.st_size,
            total_events=0,
            dropped_events=0,
            compacted_events=0,
            estimated_saved_bytes=0,
            kept_events=None,
            ok=False,
            error="not a regular file",
        )

    total = 0
    dropped = 0
    compacted = 0
    estimated_saved_bytes = 0
    kept: list[dict] | None = None if counts_only else []

    try:
        for event in _iter_events(path):
            if not isinstance(event, dict):
                return RepairPlan(
                    path=path,
                    original_size=st.st_size,
                    total_events=total,
                    dropped_events=dropped,
                    compacted_events=compacted,
                    estimated_saved_bytes=estimated_saved_bytes,
                    kept_events=None,
                    ok=False,
                    error=f"non-dict event at index {total}",
                )

            total += 1
            transformed = compact_event_for_session(event, mode="slim")
            original_len = len(json.dumps(event, ensure_ascii=False))

            if transformed is None:
                dropped += 1
                estimated_saved_bytes += original_len
                continue

            transformed_len = len(json.dumps(transformed, ensure_ascii=False))
            if transformed != event:
                compacted += 1
                estimated_saved_bytes += max(0, original_len - transformed_len)

            if kept is not None:
                kept.append(transformed)
    except SessionRepairError as e:
        return RepairPlan(
            path=path,
            original_size=st.st_size,
            total_events=total,
            dropped_events=dropped,
            compacted_events=compacted,
            estimated_saved_bytes=estimated_saved_bytes,
            kept_events=None,
            ok=False,
            error=str(e),
        )

    return RepairPlan(
        path=path,
        original_size=st.st_size,
        total_events=total,
        dropped_events=dropped,
        compacted_events=compacted,
        estimated_saved_bytes=estimated_saved_bytes,
        kept_events=kept,
        ok=True,
        error=None,
    )


# ---------- repair primitive --------------------------------------------------


def apply_repair(plan: RepairPlan, *, keep_backup: bool = False) -> None:
    """Rewrite *plan.path* with dropped/compacted events applied.

    Raises SessionRepairError on any failure; on failure the original file is
    guaranteed to be byte-for-byte unchanged and no temp file is left behind.
    """
    if not plan.ok or plan.kept_events is None:
        raise SessionRepairError(
            f"{plan.path}: cannot apply repair to a failed scan: {plan.error}"
        )

    if not plan.needs_repair:
        # Nothing to do — preserve original exactly (mtime included).
        return

    path = plan.path
    tmp = path.with_name(path.name + ".repair.tmp")
    expected_count = plan.total_events - plan.dropped_events

    try:
        with open(tmp, "w") as f:
            json.dump(plan.kept_events, f, indent=2)
            f.write("\n")
            f.flush()
            os.fsync(f.fileno())

        # Validate the temp file loads cleanly and has the right shape/count.
        with open(tmp) as f:
            reloaded = json.load(f)
        if not isinstance(reloaded, list):
            raise SessionRepairError(
                f"{path}: temp file is not a JSON list"
            )
        if len(reloaded) != expected_count:
            raise SessionRepairError(
                f"{path}: event count mismatch "
                f"(expected {expected_count}, got {len(reloaded)})"
            )
        # Any dropped type that snuck back in is a hard failure.
        for i, ev in enumerate(reloaded):
            if not isinstance(ev, dict):
                raise SessionRepairError(
                    f"{path}: non-dict event at index {i} in temp file"
                )
            if is_dropped_event(ev):
                raise SessionRepairError(
                    f"{path}: dropped-type event leaked into temp file at index {i}"
                )

        if keep_backup:
            backup = path.with_name(path.name + ".bak")
            os.replace(path, backup)

        os.replace(tmp, path)
    except BaseException:
        try:
            tmp.unlink()
        except FileNotFoundError:
            pass
        raise


# ---------- discovery ---------------------------------------------------------


def discover_session_files(abreka_dir: Path) -> list[Path]:
    """Return all pi_session*.json files under abreka_dir/experiments/*/artifacts/."""
    experiments = abreka_dir / "experiments"
    if not experiments.is_dir():
        return []
    found: list[Path] = []
    for exp in sorted(experiments.iterdir()):
        artifacts = exp / "artifacts"
        if not artifacts.is_dir():
            continue
        for p in sorted(artifacts.glob("pi_session*.json")):
            if p.is_file():
                found.append(p)
    return found
