"""Tests for WSChannel — WS Proxy as I/O adapter channel.

Uses a real aiohttp server on a random port with mocked DB (ChatDB) and
mocked Orchestrator.handle_message() to yield predetermined events.
Each test is independent.
"""

from __future__ import annotations

import asyncio
import json
import base64
import uuid
from typing import Any, AsyncIterator
from unittest.mock import AsyncMock, MagicMock, patch

import aiohttp
import aiohttp.web as web
import pytest

from march.core.orchestrator import (
    Cancelled,
    Error,
    FinalResponse,
    Orchestrator,
    TextDelta,
    ToolProgress,
)
from march.channels.ws_channel import (
    WSChannel,
    ChatDB,
    _StreamBuffer,
    _WSConn,
    _try_send,
)


# ═══════════════════════════════════════════════════════════════════════════════
# Helpers / Fixtures
# ═══════════════════════════════════════════════════════════════════════════════


class FakeChatDB:
    """In-memory ChatDB replacement for testing."""

    def __init__(self) -> None:
        self._sessions: dict[str, dict] = {}
        self._messages: dict[str, list[dict]] = {}
        self._store = MagicMock()
        self._store.db_path = "/tmp/fake.db"

    async def initialize(self) -> None:
        pass

    async def close(self) -> None:
        pass

    async def list_sessions(self) -> list[dict]:
        return list(self._sessions.values())

    async def create_session(self, name: str, description: str = "") -> dict:
        sid = str(uuid.uuid4())
        sess = {
            "id": sid,
            "name": name,
            "description": description,
            "created_at": "2026-01-01T00:00:00+00:00",
        }
        self._sessions[sid] = sess
        self._messages[sid] = []
        return sess

    async def delete_session(self, session_id: str) -> bool:
        if session_id in self._sessions:
            del self._sessions[session_id]
            self._messages.pop(session_id, None)
            return True
        return False

    async def rename_session(self, session_id: str, name: str) -> bool:
        if session_id in self._sessions:
            self._sessions[session_id]["name"] = name
            return True
        return False

    async def get_history(self, session_id: str) -> dict | None:
        if session_id not in self._sessions:
            return None
        return {
            "session": {
                "id": session_id,
                "name": self._sessions[session_id]["name"],
                "description": self._sessions[session_id].get("description", ""),
                "rolling_summary": "",
                "created_at": self._sessions[session_id]["created_at"],
                "last_active": self._sessions[session_id]["created_at"],
                "is_active": True,
            },
            "messages": self._messages.get(session_id, []),
        }

    async def session_exists(self, session_id: str) -> bool:
        return session_id in self._sessions

    async def save_message(self, session_id: str, role: str, content: str, **kw) -> str:
        msg = {"id": str(uuid.uuid4()), "role": role, "content": content}
        self._messages.setdefault(session_id, []).append(msg)
        return msg["id"]

    async def clear_session_messages(self, session_id: str) -> None:
        self._messages.pop(session_id, None)

    async def get_rolling_summary(self, session_id: str) -> str:
        return ""

    async def update_rolling_summary(self, session_id: str, summary: str) -> None:
        pass

    async def get_message_count(self, session_id: str) -> int:
        return len(self._messages.get(session_id, []))

    async def get_recent_messages(self, session_id: str, limit: int = 20) -> list[dict]:
        return self._messages.get(session_id, [])[-limit:]


def _make_plugin_and_app(
    fake_db: FakeChatDB,
    orchestrator_side_effect=None,
) -> tuple[WSChannel, web.Application]:
    """Build a WSChannel wired to a fake DB and mock orchestrator.

    Returns (channel, aiohttp_app) — caller starts the server.
    """
    plugin = WSChannel()
    plugin._db = fake_db
    plugin._agent = MagicMock()
    plugin._app_ref = MagicMock()
    plugin._max_msg_size = 20 * 1024 * 1024
    plugin._buffer_seconds = 0.0  # no drain delay in tests
    plugin._max_queue_size = 100
    plugin._max_image_dim = 512
    plugin._image_quality = 85
    plugin._max_upload_bytes = 20 * 1024 * 1024
    plugin._summary_max_tokens = 500
    plugin._summary_chunk_size = 4000

    # Mock orchestrator
    mock_orch = MagicMock(spec=Orchestrator)
    mock_orch.evict_session = MagicMock()
    if orchestrator_side_effect is not None:
        mock_orch.handle_message = orchestrator_side_effect
    else:
        # Default: yield a simple response
        async def _default_handle(*args, **kwargs):
            yield TextDelta(delta="Hello")
            yield FinalResponse(content="Hello", total_tokens=10, total_cost=0.001)
        mock_orch.handle_message = _default_handle
    plugin._orchestrator = mock_orch

    # Build the aiohttp app with routes (same as on_start but without server bind)
    webapp = web.Application()
    webapp["plugin"] = plugin
    webapp.router.add_get("/health", plugin._handle_health)
    webapp.router.add_get("/sessions", plugin._handle_list_sessions)
    webapp.router.add_post("/sessions", plugin._handle_create_session)
    webapp.router.add_delete("/sessions/{session_id}", plugin._handle_delete_session)
    webapp.router.add_put("/sessions/{session_id}", plugin._handle_rename_session)
    webapp.router.add_get("/sessions/{session_id}/history", plugin._handle_get_history)
    webapp.router.add_get("/ws/{session_id}", plugin._handle_ws)

    return plugin, webapp


@pytest.fixture
async def server_and_client():
    """Fixture that yields (plugin, fake_db, base_url, client_session).

    Starts a real aiohttp server on a random port.
    The orchestrator mock can be replaced per-test via plugin._orchestrator.
    """
    fake_db = FakeChatDB()
    plugin, webapp = _make_plugin_and_app(fake_db)

    runner = web.AppRunner(webapp)
    await runner.setup()
    site = web.TCPSite(runner, "127.0.0.1", 0)  # port 0 = random
    await site.start()

    # Extract the actual port
    sockets = site._server.sockets  # type: ignore[union-attr]
    port = sockets[0].getsockname()[1]
    base_url = f"http://127.0.0.1:{port}"

    session = aiohttp.ClientSession()

    yield plugin, fake_db, base_url, session

    await session.close()
    await site.stop()
    await runner.cleanup()


async def _create_test_session(fake_db: FakeChatDB, name: str = "Test") -> str:
    """Helper: create a session in the fake DB and return its ID."""
    sess = await fake_db.create_session(name)
    return sess["id"]


# ═══════════════════════════════════════════════════════════════════════════════
# TestWSProxyREST
# ═══════════════════════════════════════════════════════════════════════════════


class TestWSProxyREST:
    """REST endpoint tests."""

    async def test_health_endpoint(self, server_and_client):
        plugin, fake_db, base_url, client = server_and_client
        resp = await client.get(f"{base_url}/health")
        assert resp.status == 200
        data = await resp.json()
        assert data["status"] == "ok"
        assert "agent" in data

    async def test_list_sessions(self, server_and_client):
        plugin, fake_db, base_url, client = server_and_client

        # Bypass the sync fallback — patch _list_sessions_sync to use our fake_db
        sessions_data = await fake_db.list_sessions()
        plugin._list_sessions_sync = MagicMock(return_value=sessions_data)

        resp = await client.get(f"{base_url}/sessions")
        assert resp.status == 200
        data = await resp.json()
        assert "sessions" in data
        assert isinstance(data["sessions"], list)

    async def test_list_sessions_with_data(self, server_and_client):
        plugin, fake_db, base_url, client = server_and_client

        await fake_db.create_session("Chat 1")
        await fake_db.create_session("Chat 2")

        sessions_list = await fake_db.list_sessions()
        plugin._list_sessions_sync = MagicMock(return_value=sessions_list)

        resp = await client.get(f"{base_url}/sessions")
        assert resp.status == 200
        data = await resp.json()
        assert len(data["sessions"]) == 2

    async def test_create_session(self, server_and_client):
        plugin, fake_db, base_url, client = server_and_client

        # Patch _create_session_sync to use our fake_db
        async def _async_create(name, desc=""):
            return await fake_db.create_session(name, desc)

        def _sync_create(name, desc=""):
            import asyncio as _aio
            loop = _aio.new_event_loop()
            try:
                return loop.run_until_complete(_async_create(name, desc))
            finally:
                loop.close()

        plugin._create_session_sync = _sync_create

        resp = await client.post(
            f"{base_url}/sessions",
            json={"name": "My Chat", "description": "A test chat"},
        )
        assert resp.status == 201
        data = await resp.json()
        assert "id" in data
        assert data["name"] == "My Chat"
        assert data["description"] == "A test chat"

    async def test_delete_session(self, server_and_client):
        plugin, fake_db, base_url, client = server_and_client

        sid = await _create_test_session(fake_db, "To Delete")

        resp = await client.delete(f"{base_url}/sessions/{sid}")
        assert resp.status == 200
        data = await resp.json()
        assert data["deleted"] is True

        # Verify it's gone
        assert not await fake_db.session_exists(sid)

    async def test_delete_session_not_found(self, server_and_client):
        plugin, fake_db, base_url, client = server_and_client

        resp = await client.delete(f"{base_url}/sessions/nonexistent-id")
        assert resp.status == 404

    async def test_get_history(self, server_and_client):
        plugin, fake_db, base_url, client = server_and_client

        sid = await _create_test_session(fake_db, "History Test")
        await fake_db.save_message(sid, "user", "Hello")
        await fake_db.save_message(sid, "assistant", "Hi there!")

        resp = await client.get(f"{base_url}/sessions/{sid}/history")
        assert resp.status == 200
        data = await resp.json()
        assert "session" in data
        assert "messages" in data
        assert data["session"]["id"] == sid
        assert len(data["messages"]) == 2
        assert data["messages"][0]["role"] == "user"
        assert data["messages"][1]["role"] == "assistant"

    async def test_get_history_not_found(self, server_and_client):
        plugin, fake_db, base_url, client = server_and_client

        resp = await client.get(f"{base_url}/sessions/nonexistent/history")
        assert resp.status == 404


# ═══════════════════════════════════════════════════════════════════════════════
# TestWSProxyWebSocket
# ═══════════════════════════════════════════════════════════════════════════════


class TestWSProxyWebSocket:
    """WebSocket streaming tests."""

    async def test_ws_text_message(self, server_and_client):
        """Send a text message → receive stream.start, stream.delta, stream.end."""
        plugin, fake_db, base_url, client = server_and_client
        sid = await _create_test_session(fake_db)

        async def mock_handle(session_id, content, source, cancel_event=None):
            yield TextDelta(delta="Hello ")
            yield TextDelta(delta="world!")
            yield FinalResponse(content="Hello world!", total_tokens=50, total_cost=0.01)

        plugin._orchestrator.handle_message = mock_handle

        ws_url = f"{base_url}/ws/{sid}"
        async with client.ws_connect(ws_url) as ws:
            await ws.send_json({"type": "message", "content": "Hi"})

            events = []
            while True:
                msg = await asyncio.wait_for(ws.receive_json(), timeout=5.0)
                events.append(msg)
                if msg.get("type") in ("stream.end", "error"):
                    break

            types = [e["type"] for e in events]
            assert "stream.start" in types
            assert "stream.delta" in types
            assert "stream.end" in types

            # Verify deltas
            deltas = [e for e in events if e["type"] == "stream.delta"]
            assert len(deltas) == 2
            assert deltas[0]["content"] == "Hello "
            assert deltas[1]["content"] == "world!"

    async def test_ws_stop_command(self, server_and_client):
        """Send 'stop' during stream → stream.cancelled.

        The WS handler processes messages sequentially, so we simulate
        cancellation by setting cancel_event externally (as the orchestrator
        would see it) while the stream is in progress.
        """
        plugin, fake_db, base_url, client = server_and_client
        sid = await _create_test_session(fake_db)

        # We'll capture the cancel_event passed to handle_message and set it
        # from a background task to simulate the stop being processed.
        captured_cancel = {}
        mock_started = asyncio.Event()

        async def mock_handle(session_id, content, source, cancel_event=None):
            captured_cancel["event"] = cancel_event
            yield TextDelta(delta="Starting...")
            mock_started.set()
            # Poll for cancel
            for _ in range(200):
                if cancel_event and cancel_event.is_set():
                    yield Cancelled(partial_content="Starting...")
                    return
                await asyncio.sleep(0.05)
            yield FinalResponse(content="Starting...", total_tokens=10, total_cost=0.001)

        plugin._orchestrator.handle_message = mock_handle

        ws_url = f"{base_url}/ws/{sid}"
        async with client.ws_connect(ws_url) as ws:
            await ws.send_json({"type": "message", "content": "Tell me a story"})

            # Read stream.start
            msg = await asyncio.wait_for(ws.receive_json(), timeout=5.0)
            assert msg["type"] == "stream.start"

            # Read the delta
            msg = await asyncio.wait_for(ws.receive_json(), timeout=5.0)
            assert msg["type"] == "stream.delta"

            # Ensure mock is in its polling phase
            await asyncio.wait_for(mock_started.wait(), timeout=5.0)

            # Simulate stop: set the cancel_event directly (as the WS handler
            # would do if it could process the stop message concurrently)
            captured_cancel["event"].set()

            # Collect remaining events — should get stream.end with cancelled=True
            events = []
            while True:
                msg = await asyncio.wait_for(ws.receive_json(), timeout=5.0)
                events.append(msg)
                if msg.get("type") in ("stream.end", "error"):
                    break

            types = [e["type"] for e in events]
            assert "stream.end" in types
            end_event = [e for e in events if e["type"] == "stream.end"][0]
            assert end_event.get("cancelled") is True

    async def test_ws_tool_progress(self, server_and_client):
        """Agent uses tool → tool.progress event received."""
        plugin, fake_db, base_url, client = server_and_client
        sid = await _create_test_session(fake_db)

        async def mock_handle(session_id, content, source, cancel_event=None):
            yield ToolProgress(name="web_search", status="started", summary="Searching...")
            yield ToolProgress(
                name="web_search", status="complete",
                summary="Found 3 results", duration_ms=150.0,
            )
            yield TextDelta(delta="Here are the results.")
            yield FinalResponse(content="Here are the results.", total_tokens=100, total_cost=0.02)

        plugin._orchestrator.handle_message = mock_handle

        ws_url = f"{base_url}/ws/{sid}"
        async with client.ws_connect(ws_url) as ws:
            await ws.send_json({"type": "message", "content": "Search for cats"})

            events = []
            while True:
                msg = await asyncio.wait_for(ws.receive_json(), timeout=5.0)
                events.append(msg)
                if msg.get("type") in ("stream.end", "error"):
                    break

            tool_events = [e for e in events if e["type"] == "tool.progress"]
            assert len(tool_events) == 2
            assert tool_events[0]["name"] == "web_search"
            assert tool_events[0]["status"] == "started"
            assert tool_events[1]["status"] == "complete"
            assert tool_events[1]["duration_ms"] == 150.0

    async def test_ws_token_reporting(self, server_and_client):
        """stream.end has correct usage (total_tokens, cost)."""
        plugin, fake_db, base_url, client = server_and_client
        sid = await _create_test_session(fake_db)

        async def mock_handle(session_id, content, source, cancel_event=None):
            yield TextDelta(delta="Response text")
            yield FinalResponse(
                content="Response text",
                total_tokens=1234,
                total_cost=0.0567,
            )

        plugin._orchestrator.handle_message = mock_handle

        ws_url = f"{base_url}/ws/{sid}"
        async with client.ws_connect(ws_url) as ws:
            await ws.send_json({"type": "message", "content": "Hello"})

            events = []
            while True:
                msg = await asyncio.wait_for(ws.receive_json(), timeout=5.0)
                events.append(msg)
                if msg.get("type") in ("stream.end", "error"):
                    break

            end_event = [e for e in events if e["type"] == "stream.end"][0]
            assert "usage" in end_event
            assert end_event["usage"]["total_tokens"] == 1234
            assert abs(end_event["usage"]["cost"] - 0.0567) < 1e-6


# ═══════════════════════════════════════════════════════════════════════════════
# TestWSProxyMultiModal
# ═══════════════════════════════════════════════════════════════════════════════


class TestWSProxyMultiModal:
    """Multimodal attachment tests (image, voice, PDF)."""

    async def test_ws_image_upload(self, server_and_client):
        """Send image attachment → processed and passed to orchestrator."""
        plugin, fake_db, base_url, client = server_and_client
        sid = await _create_test_session(fake_db)

        received_content = {}

        async def mock_handle(session_id, content, source, cancel_event=None):
            received_content["value"] = content
            yield TextDelta(delta="Nice image!")
            yield FinalResponse(content="Nice image!", total_tokens=20, total_cost=0.001)

        plugin._orchestrator.handle_message = mock_handle

        # Create a tiny valid JPEG (smallest possible)
        # 1x1 white pixel JPEG
        jpeg_bytes = base64.b64decode(
            "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkS"
            "Ew8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/2wBDAQkJ"
            "CQwLDBgNDRgyIRwhMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIy"
            "MjIyMjIyMjIyMjIyMjL/wAARCAABAAEDASIAAhEBAxEB/8QAHwAAAQUBAQEBAQEA"
            "AAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIh"
            "MUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6"
            "Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZ"
            "mqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx"
            "8vP09fb3+Pn6/8QAHwEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREA"
            "AgECBAQDBAcFBAQAAQJ3AAECAxEEBSExBhJBUQdhcRMiMoEIFEKRobHBCSMzUvAV"
            "YnLRChYkNOEl8RcYI4Q/RFhHRUYnJCk6OTtBRUpHTFFWV1hZWmNkZWZnaGlqc3R1"
            "dnd4eXqCg4SFhoeIiYqSk5SVlpeYmZqio6Slpqeoqaqys7S1tre4ubrCw8TFxsfI"
            "ycrS09TV1tfY2dri4+Tl5ufo6ery8/T19vf4+fr/2gAMAwEAAhEDEQA/AP0poA//2Q=="
        )

        with patch("march.channels.ws_channel._process_attachment") as mock_proc:
            # Return multimodal content (image block + text)
            mock_proc.return_value = [
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/jpeg",
                        "data": base64.b64encode(jpeg_bytes).decode(),
                    },
                },
                {
                    "type": "text",
                    "text": "[User attached image: test.jpg (1KB)]",
                },
            ]

            ws_url = f"{base_url}/ws/{sid}"
            async with client.ws_connect(ws_url) as ws:
                await ws.send_json({
                    "type": "attachment",
                    "filename": "test.jpg",
                    "mime_type": "image/jpeg",
                    "data": base64.b64encode(jpeg_bytes).decode(),
                })

                events = []
                while True:
                    msg = await asyncio.wait_for(ws.receive_json(), timeout=5.0)
                    events.append(msg)
                    if msg.get("type") in ("stream.end", "error"):
                        break

                types = [e["type"] for e in events]
                # Should have image.preview and stream events
                assert "stream.end" in types

                # Verify orchestrator received multimodal content
                assert received_content.get("value") is not None
                content = received_content["value"]
                assert isinstance(content, list)
                assert any(b.get("type") == "image" for b in content)

    async def test_ws_voice_upload(self, server_and_client):
        """Send voice attachment → transcribed and passed to orchestrator."""
        plugin, fake_db, base_url, client = server_and_client
        sid = await _create_test_session(fake_db)

        received_content = {}

        async def mock_handle(session_id, content, source, cancel_event=None):
            received_content["value"] = content
            yield TextDelta(delta="Got it!")
            yield FinalResponse(content="Got it!", total_tokens=15, total_cost=0.001)

        plugin._orchestrator.handle_message = mock_handle

        # Mock the voice transcription tool
        mock_tool_result = MagicMock()
        mock_tool_result.is_error = False
        mock_tool_result.content = "Hello, this is a voice message"

        mock_tools = MagicMock()
        mock_tools.execute = AsyncMock(return_value=mock_tool_result)
        plugin._agent.tools = mock_tools

        fake_audio = base64.b64encode(b"\x00" * 100).decode()

        ws_url = f"{base_url}/ws/{sid}"
        async with client.ws_connect(ws_url) as ws:
            await ws.send_json({
                "type": "voice",
                "mime_type": "audio/webm",
                "data": fake_audio,
            })

            events = []
            while True:
                msg = await asyncio.wait_for(ws.receive_json(), timeout=5.0)
                events.append(msg)
                if msg.get("type") in ("stream.end", "error"):
                    break

            types = [e["type"] for e in events]
            assert "voice.transcribed" in types
            assert "stream.end" in types

            # Verify the transcribed text was sent to orchestrator
            assert received_content["value"] == "Hello, this is a voice message"

            # Verify voice.transcribed event has the text
            vt = [e for e in events if e["type"] == "voice.transcribed"][0]
            assert vt["text"] == "Hello, this is a voice message"

    async def test_ws_pdf_upload(self, server_and_client):
        """Send PDF → extracted and passed to orchestrator."""
        plugin, fake_db, base_url, client = server_and_client
        sid = await _create_test_session(fake_db)

        received_content = {}

        async def mock_handle(session_id, content, source, cancel_event=None):
            received_content["value"] = content
            yield TextDelta(delta="PDF processed.")
            yield FinalResponse(content="PDF processed.", total_tokens=30, total_cost=0.002)

        plugin._orchestrator.handle_message = mock_handle

        fake_pdf = base64.b64encode(b"%PDF-1.4 fake content").decode()

        with patch("march.channels.ws_channel._process_attachment") as mock_proc:
            mock_proc.return_value = (
                "[media attached: /tmp/test.pdf (PDF, 3 pages, 10KB)]\n\n"
                "Summary: This document discusses testing strategies."
            )

            ws_url = f"{base_url}/ws/{sid}"
            async with client.ws_connect(ws_url) as ws:
                await ws.send_json({
                    "type": "attachment",
                    "filename": "report.pdf",
                    "mime_type": "application/pdf",
                    "data": fake_pdf,
                })

                events = []
                while True:
                    msg = await asyncio.wait_for(ws.receive_json(), timeout=5.0)
                    events.append(msg)
                    if msg.get("type") in ("stream.end", "error"):
                        break

                types = [e["type"] for e in events]
                assert "stream.end" in types

                # Verify orchestrator received the PDF summary text
                val = received_content["value"]
                assert isinstance(val, str)
                assert "PDF" in val


# ═══════════════════════════════════════════════════════════════════════════════
# TestWSProxyConnection
# ═══════════════════════════════════════════════════════════════════════════════


class TestWSProxyConnection:
    """Connection management tests: takeover, reconnect, continuity."""

    async def test_ws_takeover(self, server_and_client):
        """Second client connects → first gets session.takeover."""
        plugin, fake_db, base_url, client = server_and_client
        sid = await _create_test_session(fake_db)

        ws_url = f"{base_url}/ws/{sid}"

        # First client connects
        ws1 = await client.ws_connect(ws_url)

        # Second client connects — first should get takeover
        ws2 = await client.ws_connect(ws_url)

        # First client should receive session.takeover then close
        msg = await asyncio.wait_for(ws1.receive(), timeout=5.0)
        if msg.type == aiohttp.WSMsgType.TEXT:
            data = json.loads(msg.data)
            assert data["type"] == "session.takeover"
            assert "message" in data
        elif msg.type == aiohttp.WSMsgType.CLOSE:
            # Connection was closed — that's also acceptable after takeover
            pass

        await ws1.close()
        await ws2.close()

    async def test_ws_reconnect_during_stream(self, server_and_client):
        """Disconnect mid-stream, reconnect → stream.active with collected content.

        Uses the _StreamBuffer directly to verify reconnect recovery behavior,
        since aiohttp's ws.close() blocks until the handler finishes.
        """
        plugin, fake_db, base_url, client = server_and_client
        sid = await _create_test_session(fake_db)

        async def mock_handle(session_id, content, source, cancel_event=None):
            yield TextDelta(delta="Part 1. ")
            yield TextDelta(delta="Part 2. ")
            yield FinalResponse(content="Part 1. Part 2.", total_tokens=50, total_cost=0.01)

        plugin._orchestrator.handle_message = mock_handle

        ws_url = f"{base_url}/ws/{sid}"

        # First client: send a message and let the stream complete
        async with client.ws_connect(ws_url) as ws1:
            await ws1.send_json({"type": "message", "content": "Hello"})

            events1 = []
            while True:
                msg = await asyncio.wait_for(ws1.receive_json(), timeout=5.0)
                events1.append(msg)
                if msg.get("type") in ("stream.end", "error"):
                    break

        # After disconnect, the stream buffer should have collected content
        # and be marked as done. On reconnect, client gets stream.catchup.
        async with client.ws_connect(ws_url) as ws2:
            msg = await asyncio.wait_for(ws2.receive_json(), timeout=5.0)
            assert msg["type"] == "stream.catchup"
            assert msg["done"] is True
            assert "Part 1." in msg["content"]
            assert "Part 2." in msg["content"]

    async def test_ws_session_continuity(self, server_and_client):
        """Connect, send msg, disconnect, reconnect → history preserved."""
        plugin, fake_db, base_url, client = server_and_client
        sid = await _create_test_session(fake_db)

        call_count = {"n": 0}

        async def mock_handle(session_id, content, source, cancel_event=None):
            call_count["n"] += 1
            if call_count["n"] == 1:
                yield TextDelta(delta="First response")
                yield FinalResponse(content="First response", total_tokens=10, total_cost=0.001)
            else:
                yield TextDelta(delta="Second response")
                yield FinalResponse(content="Second response", total_tokens=10, total_cost=0.001)

        plugin._orchestrator.handle_message = mock_handle

        ws_url = f"{base_url}/ws/{sid}"

        # First connection: send a message
        async with client.ws_connect(ws_url) as ws1:
            await ws1.send_json({"type": "message", "content": "Hello"})
            events1 = []
            while True:
                msg = await asyncio.wait_for(ws1.receive_json(), timeout=5.0)
                events1.append(msg)
                if msg.get("type") in ("stream.end", "error"):
                    break

        # Save a message to fake DB to simulate persistence
        await fake_db.save_message(sid, "user", "Hello")
        await fake_db.save_message(sid, "assistant", "First response")

        # Verify history is preserved via REST
        resp = await client.get(f"{base_url}/sessions/{sid}/history")
        assert resp.status == 200
        data = await resp.json()
        assert len(data["messages"]) == 2

        # Reconnect and send another message
        async with client.ws_connect(ws_url) as ws2:
            await ws2.send_json({"type": "message", "content": "Follow up"})
            events2 = []
            while True:
                msg = await asyncio.wait_for(ws2.receive_json(), timeout=5.0)
                events2.append(msg)
                if msg.get("type") in ("stream.end", "error"):
                    break

            types2 = [e["type"] for e in events2]
            assert "stream.end" in types2

        # Verify both exchanges are in history
        await fake_db.save_message(sid, "user", "Follow up")
        await fake_db.save_message(sid, "assistant", "Second response")

        resp = await client.get(f"{base_url}/sessions/{sid}/history")
        data = await resp.json()
        assert len(data["messages"]) == 4
        assert data["messages"][0]["content"] == "Hello"
        assert data["messages"][3]["content"] == "Second response"
