"""Tests for abreka.session_repair."""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from abreka.session_repair import (
    RepairPlan,
    apply_repair,
    discover_session_files,
    scan_session_file,
)


def _write_session(path: Path, events: list[dict]) -> None:
    """Write a session file in the exact format PiRpcClient.save_session uses."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        json.dump(events, f, indent=2)
        f.write("\n")


def _delta(payload: str = "a") -> dict:
    """Build a droppable message_update/toolcall_delta event."""
    return {
        "type": "message_update",
        "assistantMessageEvent": {"type": "toolcall_delta", "delta": payload},
    }


def _ame_end() -> dict:
    return {
        "type": "message_update",
        "assistantMessageEvent": {"type": "toolcall_end"},
    }


def _snapshot_bloat() -> dict:
    snapshot = {
        "role": "assistant",
        "content": [{"type": "thinking", "thinking": "plan"}],
    }
    return {
        "type": "message_update",
        "message": snapshot,
        "assistantMessageEvent": {
            "type": "thinking_delta",
            "delta": "x",
            "partial": snapshot,
        },
    }


# ---------- scan_session_file -------------------------------------------------


class TestScanSessionFile:
    def test_scan_file_with_deltas(self, tmp_path):
        events = [
            {"type": "message_start", "message": {"role": "user"}},
            _delta("a"),
            _delta("ab"),
            _ame_end(),
            {"type": "message_end", "message": {}},
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)

        plan = scan_session_file(f)

        assert plan.ok is True
        assert plan.error is None
        assert plan.path == f
        assert plan.original_size == f.stat().st_size
        assert plan.total_events == 5
        assert plan.dropped_events == 2
        assert plan.kept_events is not None
        # Three kept events: message_start, ame:toolcall_end, message_end.
        assert len(plan.kept_events) == 3
        assert plan.kept_events[0]["type"] == "message_start"
        assert plan.kept_events[1]["assistantMessageEvent"]["type"] == "toolcall_end"
        assert plan.kept_events[2]["type"] == "message_end"

    def test_scan_file_without_deltas(self, tmp_path):
        events = [
            {"type": "message_start", "message": {"role": "user"}},
            _ame_end(),
            {"type": "message_end", "message": {}},
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)

        plan = scan_session_file(f)
        assert plan.ok is True
        assert plan.total_events == 3
        assert plan.dropped_events == 0
        assert plan.compacted_events == 0
        assert len(plan.kept_events) == 3

    def test_scan_marks_snapshot_payloads_as_compactable(self, tmp_path):
        events = [
            {"type": "message_start", "message": {"role": "assistant"}},
            _snapshot_bloat(),
            {"type": "message_end", "message": {"role": "assistant"}},
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)

        plan = scan_session_file(f)

        assert plan.ok is True
        assert plan.total_events == 3
        assert plan.dropped_events == 0
        assert plan.compacted_events == 1
        assert plan.estimated_saved_bytes > 0
        assert plan.kept_events is not None
        update = plan.kept_events[1]
        assert "message" not in update
        assert "partial" not in update["assistantMessageEvent"]

    def test_scan_empty_list(self, tmp_path):
        f = tmp_path / "s.json"
        _write_session(f, [])
        plan = scan_session_file(f)
        assert plan.ok is True
        assert plan.total_events == 0
        assert plan.dropped_events == 0
        assert plan.kept_events == []

    def test_scan_rejects_non_array_json(self, tmp_path):
        f = tmp_path / "s.json"
        f.write_text('{"not": "a list"}\n')
        plan = scan_session_file(f)
        assert plan.ok is False
        assert plan.error is not None
        assert plan.kept_events is None

    def test_scan_rejects_non_json(self, tmp_path):
        f = tmp_path / "s.json"
        f.write_text("hello world\n")
        plan = scan_session_file(f)
        assert plan.ok is False
        assert plan.error is not None

    def test_scan_rejects_directory(self, tmp_path):
        d = tmp_path / "notafile"
        d.mkdir()
        plan = scan_session_file(d)
        assert plan.ok is False

    def test_counts_only_does_not_retain_events(self, tmp_path):
        events = [
            {"type": "message_start", "message": {"role": "user"}},
            _delta("a"),
            _ame_end(),
            {"type": "message_end", "message": {}},
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)
        plan = scan_session_file(f, counts_only=True)
        assert plan.ok is True
        assert plan.total_events == 4
        assert plan.dropped_events == 1
        assert plan.compacted_events == 0
        assert plan.kept_events is None  # memory-saving mode

    def test_counts_only_plan_cannot_be_applied(self, tmp_path):
        events = [
            _delta("a"),
            {"type": "message_end", "message": {}},
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)
        plan = scan_session_file(f, counts_only=True)
        with pytest.raises(Exception):
            apply_repair(plan)
        # Original must be byte-identical.
        reloaded = json.loads(f.read_text())
        assert len(reloaded) == 2

    def test_scan_handles_large_single_event(self, tmp_path):
        big_text = "x" * 2_000_000  # 2 MB of content
        events = [
            {"type": "message_start", "message": {"role": "user"}},
            {"type": "tool_execution_end", "result": {"content": [{"type": "text", "text": big_text}]}},
            _delta("junk"),
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)

        plan = scan_session_file(f)
        assert plan.ok is True
        assert plan.total_events == 3
        assert plan.dropped_events == 1
        assert plan.compacted_events == 0
        assert len(plan.kept_events) == 2
        assert plan.kept_events[1]["result"]["content"][0]["text"] == big_text


# ---------- apply_repair ------------------------------------------------------


class TestApplyRepair:
    def test_atomic_rewrite(self, tmp_path):
        events = [
            {"type": "message_start", "message": {"role": "user"}},
            _delta("a"),
            _ame_end(),
            {"type": "message_end", "message": {}},
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)
        original_size = f.stat().st_size

        plan = scan_session_file(f)
        apply_repair(plan)

        reloaded = json.loads(f.read_text())
        assert len(reloaded) == 3
        assert reloaded[0]["type"] == "message_start"
        assert reloaded[1]["assistantMessageEvent"]["type"] == "toolcall_end"
        assert reloaded[2]["type"] == "message_end"
        assert f.stat().st_size < original_size
        assert not (tmp_path / "s.json.repair.tmp").exists()
        assert not (tmp_path / "s.json.bak").exists()

    def test_noop_when_nothing_to_drop(self, tmp_path):
        events = [
            {"type": "message_start", "message": {"role": "user"}},
            {"type": "message_end", "message": {}},
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)
        original_bytes = f.read_bytes()
        original_mtime = f.stat().st_mtime_ns

        plan = scan_session_file(f)
        apply_repair(plan)

        assert f.read_bytes() == original_bytes
        assert f.stat().st_mtime_ns == original_mtime

    def test_rewrites_file_when_only_compaction_is_needed(self, tmp_path):
        events = [
            {"type": "message_start", "message": {"role": "assistant"}},
            _snapshot_bloat(),
            {"type": "message_end", "message": {"role": "assistant"}},
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)
        original_size = f.stat().st_size

        plan = scan_session_file(f)
        assert plan.needs_repair is True
        assert plan.dropped_events == 0
        assert plan.compacted_events == 1

        apply_repair(plan)

        reloaded = json.loads(f.read_text())
        assert len(reloaded) == 3
        assert "message" not in reloaded[1]
        assert "partial" not in reloaded[1]["assistantMessageEvent"]
        assert f.stat().st_size < original_size

    def test_keep_backup(self, tmp_path):
        events = [
            {"type": "message_start", "message": {"role": "user"}},
            _delta("x"),
            {"type": "message_end", "message": {}},
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)
        original_bytes = f.read_bytes()

        plan = scan_session_file(f)
        apply_repair(plan, keep_backup=True)

        bak = tmp_path / "s.json.bak"
        assert bak.exists()
        assert bak.read_bytes() == original_bytes
        reloaded = json.loads(f.read_text())
        assert len(reloaded) == 2

    def test_validates_before_replace(self, tmp_path, monkeypatch):
        """If the temp file somehow ends up invalid, the original must be untouched."""
        events = [
            {"type": "message_start", "message": {"role": "user"}},
            _delta("x"),
            {"type": "message_end", "message": {}},
        ]
        f = tmp_path / "s.json"
        _write_session(f, events)
        original_bytes = f.read_bytes()

        plan = scan_session_file(f)

        # Corrupt the plan's kept_events so the temp file will parse to the
        # wrong count and validation should reject it.
        plan.kept_events.append(_delta("sneak"))

        with pytest.raises(Exception):
            apply_repair(plan)

        # Original must be byte-identical
        assert f.read_bytes() == original_bytes
        assert not (tmp_path / "s.json.repair.tmp").exists()

    def test_refuses_bad_plan(self, tmp_path):
        """apply_repair should refuse plans that aren't ok."""
        f = tmp_path / "s.json"
        f.write_text("garbage")
        plan = scan_session_file(f)
        assert plan.ok is False
        with pytest.raises(Exception):
            apply_repair(plan)
        assert f.read_text() == "garbage"  # untouched


# ---------- discover_session_files --------------------------------------------


class TestDiscoverSessionFiles:
    def test_finds_session_files(self, tmp_path):
        abreka_dir = tmp_path / "abreka"
        for exp_id in ("0001", "0002"):
            artifacts = abreka_dir / "experiments" / exp_id / "artifacts"
            artifacts.mkdir(parents=True)
            (artifacts / "pi_session.json").write_text("[]\n")
            (artifacts / "pi_session_implement.json").write_text("[]\n")
            (artifacts / "run_output.txt").write_text("nope")  # not a session

        found = discover_session_files(abreka_dir)
        names = sorted(p.name for p in found)
        assert names == [
            "pi_session.json",
            "pi_session.json",
            "pi_session_implement.json",
            "pi_session_implement.json",
        ]

    def test_empty_abreka_dir(self, tmp_path):
        abreka_dir = tmp_path / "abreka"
        (abreka_dir / "experiments").mkdir(parents=True)
        assert discover_session_files(abreka_dir) == []
