"""Tests for abreka.pi_client."""

import json
import subprocess
from pathlib import Path

import pytest

from abreka.pi_client import (
    DROPPED_AME_SUBTYPES,
    PiRpcClient,
    SESSION_CAPTURE_MODES,
    is_dropped_event,
)


class TestPiRpcClientDataclass:
    def test_eof_is_instance_field(self):
        """_eof must be a dataclass instance field, not a class variable."""
        c = PiRpcClient(provider="x", model="y", cwd=Path("."))
        assert "_eof" in c.__dataclass_fields__
        assert c._eof is False

    def test_eof_not_shared_between_instances(self):
        """Each PiRpcClient instance must have its own _eof state."""
        c1 = PiRpcClient(provider="x", model="y", cwd=Path("."))
        c2 = PiRpcClient(provider="x", model="y", cwd=Path("."))
        c1._eof = True
        assert c2._eof is False

    def test_all_private_fields_are_dataclass_fields(self):
        """All _-prefixed attributes should be declared as dataclass fields."""
        c = PiRpcClient(provider="x", model="y", cwd=Path("."))
        for name in ("_process", "_all_events", "_eof"):
            assert name in c.__dataclass_fields__, f"{name} is not a dataclass field"

    def test_thinking_level_in_start_cmd(self):
        """thinking_level should default to xhigh."""
        c = PiRpcClient(provider="test", model="test", cwd=Path("."))
        assert c.thinking_level == "xhigh"

    def test_session_capture_mode_defaults_to_slim(self):
        c = PiRpcClient(provider="test", model="test", cwd=Path("."))
        assert c.session_capture_mode == "slim"
        assert SESSION_CAPTURE_MODES == frozenset({"full", "slim", "minimal"})

    def test_invalid_session_capture_mode_rejected(self):
        with pytest.raises(ValueError):
            PiRpcClient(
                provider="test",
                model="test",
                cwd=Path("."),
                session_capture_mode="not-a-mode",
            )


class TestPiRpcClientStart:
    def test_start_uses_explicit_session_path_when_provided(self, tmp_path, monkeypatch):
        session_path = tmp_path / "sessions" / "implement.jsonl"
        recorded: dict[str, object] = {}

        class DummyProcess:
            stdin = None
            stdout = None
            stderr = None

        def fake_popen(cmd, **kwargs):
            recorded["cmd"] = cmd
            recorded["cwd"] = kwargs.get("cwd")
            return DummyProcess()

        monkeypatch.setattr(subprocess, "Popen", fake_popen)

        client = PiRpcClient(
            provider="test",
            model="test",
            cwd=tmp_path,
            session_path=session_path,
        )
        client.start()

        assert session_path.parent.is_dir()
        assert recorded["cmd"][:3] == ["pi", "--mode", "rpc"]
        assert "--session" in recorded["cmd"]
        assert str(session_path) in recorded["cmd"]
        assert "--no-session" not in recorded["cmd"]
        assert recorded["cwd"] == tmp_path

def _ame(subtype: str, **extra) -> dict:
    """Build a top-level message_update event with the given ame subtype."""
    return {
        "type": "message_update",
        "assistantMessageEvent": {"type": subtype, **extra},
    }


class TestDroppedAmeSubtypes:
    """Session files should not persist cumulative-snapshot streaming events."""

    def test_toolcall_delta_is_dropped(self):
        assert "toolcall_delta" in DROPPED_AME_SUBTYPES

    def test_drop_set_is_narrow(self):
        """Pins the drop set to exactly one subtype.

        Widening this set must be a deliberate change: each type removed is
        data that downstream consumers can never recover.
        """
        assert DROPPED_AME_SUBTYPES == frozenset({"toolcall_delta"})


class TestIsDroppedEvent:
    def test_toolcall_delta_ame_is_dropped(self):
        assert is_dropped_event(_ame("toolcall_delta", delta="abc")) is True

    def test_thinking_delta_ame_is_kept(self):
        assert is_dropped_event(_ame("thinking_delta", delta="hmm")) is False

    def test_text_delta_ame_is_kept(self):
        assert is_dropped_event(_ame("text_delta", delta="hi")) is False

    def test_other_ame_subtypes_kept(self):
        for sub in (
            "toolcall_start",
            "toolcall_end",
            "thinking_start",
            "thinking_end",
            "text_start",
            "text_end",
        ):
            assert is_dropped_event(_ame(sub)) is False

    def test_top_level_message_update_without_ame_kept(self):
        assert is_dropped_event({"type": "message_update"}) is False
        assert is_dropped_event({"type": "message_update", "assistantMessageEvent": None}) is False
        assert is_dropped_event({"type": "message_update", "assistantMessageEvent": "wat"}) is False

    def test_non_message_update_kept(self):
        """Even an event literally named `toolcall_delta` at top level is kept.

        Real pi-rpc streams never emit that, but the filter is defined over
        the ame subtype specifically so any top-level type passes through.
        """
        assert is_dropped_event({"type": "toolcall_delta"}) is False
        assert is_dropped_event({"type": "message_start"}) is False
        assert is_dropped_event({"type": "tool_execution_end"}) is False


class TestRecordEvent:
    def test_record_event_drops_toolcall_delta_ame(self):
        c = PiRpcClient(provider="x", model="y", cwd=Path("."))
        c._record_event(_ame("toolcall_delta", delta="abc"))
        c._record_event({"type": "message_end", "message": {}})
        assert len(c._all_events) == 1
        assert c._all_events[0]["type"] == "message_end"

    def test_record_event_keeps_other_types(self):
        """Non-dropped events must flow through unchanged."""
        c = PiRpcClient(provider="x", model="y", cwd=Path("."))
        keepers = [
            {"type": "message_start", "message": {"role": "user"}},
            _ame("thinking_delta", delta="hmm"),
            _ame("text_delta", delta="hi"),
            _ame("toolcall_start"),
            _ame("toolcall_end"),
            {"type": "message_end", "message": {}},
            {"type": "tool_execution_start", "toolCallId": "abc"},
            {"type": "tool_execution_end", "toolCallId": "abc"},
            {"type": "turn_end", "message": {}},
            {"type": "agent_end"},
            {"type": "something_brand_new"},
        ]
        for e in keepers:
            c._record_event(e)
        assert c._all_events == keepers

    def test_record_event_preserves_order_with_mix(self):
        c = PiRpcClient(provider="x", model="y", cwd=Path("."))
        seq = [
            _ame("toolcall_start"),
            _ame("toolcall_delta", delta="a"),
            _ame("toolcall_delta", delta="ab"),
            _ame("toolcall_end"),
            _ame("toolcall_delta", delta="c"),
            {"type": "message_end"},
        ]
        for e in seq:
            c._record_event(e)
        kept_subtypes = []
        for e in c._all_events:
            if e.get("type") == "message_update":
                kept_subtypes.append(e["assistantMessageEvent"]["type"])
            else:
                kept_subtypes.append(e["type"])
        assert kept_subtypes == [
            "toolcall_start",
            "toolcall_end",
            "message_end",
        ]

    def test_slim_mode_strips_redundant_message_snapshots(self):
        c = PiRpcClient(provider="x", model="y", cwd=Path("."), session_capture_mode="slim")
        event = {
            "type": "message_update",
            "message": {"role": "assistant", "content": [{"type": "text", "text": "snapshot"}]},
            "assistantMessageEvent": {
                "type": "thinking_delta",
                "delta": "x",
                "partial": {"role": "assistant", "content": [{"type": "text", "text": "snapshot"}]},
            },
        }
        c._record_event(event)

        kept = c._all_events[0]
        assert "message" not in kept
        assert "partial" not in kept["assistantMessageEvent"]
        assert kept["assistantMessageEvent"]["delta"] == "x"

    def test_minimal_mode_drops_thinking_delta_but_keeps_text_delta(self):
        c = PiRpcClient(provider="x", model="y", cwd=Path("."), session_capture_mode="minimal")
        c._record_event(_ame("thinking_delta", delta="hmm"))
        c._record_event(_ame("text_delta", delta="hi"))
        c._record_event({"type": "message_end", "message": {"role": "assistant"}})

        assert len(c._all_events) == 2
        assert c._all_events[0]["assistantMessageEvent"]["type"] == "text_delta"
        assert c._all_events[1]["type"] == "message_end"


class TestPromptCollection:
    def _feed(self, client: PiRpcClient, events: list[dict]) -> None:
        it = iter(events)

        def _fake_send(_cmd: dict) -> None:
            return None

        def _fake_read_event() -> dict | None:
            try:
                return next(it)
            except StopIteration:
                client._eof = True
                return None

        client._send = _fake_send  # type: ignore[method-assign]
        client._read_event = _fake_read_event  # type: ignore[method-assign]

    def test_prompt_default_does_not_collect_raw_events(self):
        c = PiRpcClient(provider="x", model="y", cwd=Path("."))
        stream = [_ame("text_delta", delta="hello"), {"type": "agent_end"}]
        self._feed(c, stream)

        result = c.prompt("hi")

        assert result.text == "hello"
        assert result.events == []

    def test_prompt_can_collect_events_when_enabled(self):
        c = PiRpcClient(
            provider="x",
            model="y",
            cwd=Path("."),
            collect_result_events=True,
        )
        stream = [_ame("text_delta", delta="hello"), {"type": "agent_end"}]
        self._feed(c, stream)

        result = c.prompt("hi")

        assert result.text == "hello"
        assert len(result.events) == 2
        assert result.events[0]["assistantMessageEvent"]["type"] == "text_delta"


class TestSaveSessionRoundtrip:
    def test_save_session_excludes_dropped_events(self, tmp_path):
        c = PiRpcClient(provider="x", model="y", cwd=Path("."))
        for e in [
            {"type": "message_start", "message": {"role": "user"}},
            _ame("toolcall_delta", delta="xxx"),
            _ame("toolcall_end"),
            _ame("toolcall_delta", delta="yyy"),
            {"type": "message_end", "message": {}},
        ]:
            c._record_event(e)

        out = tmp_path / "session.json"
        c.save_session(out)

        loaded = json.loads(out.read_text())
        assert [is_dropped_event(e) for e in loaded] == [False] * len(loaded)
        assert len(loaded) == 3  # message_start, toolcall_end ame, message_end
