import asyncio
from pathlib import Path
from types import SimpleNamespace

import pytest

from beaver.engine import AgentLoop, EngineLoader
from beaver.engine.providers.base import LLMProvider, LLMResponse
from beaver.engine.providers.factory import ProviderBundle
from beaver.memory.curated.snapshot import MemorySnapshot
from beaver.memory.curated.store import MemoryStore


class _RecordingProvider(LLMProvider):
    def __init__(self, response_text: str = "done") -> None:
        super().__init__()
        self.response_text = response_text
        self.messages: list[list[dict]] = []

    async def chat(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        model: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
        thinking_enabled: bool | None = None,
    ) -> LLMResponse:
        self.messages.append(messages)
        return LLMResponse(
            content=self.response_text,
            finish_reason="stop",
            provider_name="stub",
            model="stub-model",
        )

    def get_default_model(self) -> str:
        return "stub-model"


class _FakeMemoryService:
    def __init__(self, root: Path) -> None:
        self.store = MemoryStore(root)
        self.snapshot_calls: list[dict] = []
        self.archive_calls: list[dict] = []
        self.archive_started: asyncio.Event | None = None
        self.archive_release: asyncio.Event | None = None

    def initialize(self) -> None:
        self.store.load_from_disk()

    def get_store(self) -> MemoryStore:
        return self.store

    async def capture_snapshot_for_run_async(self, **kwargs) -> MemorySnapshot:
        self.snapshot_calls.append(kwargs)
        return MemorySnapshot(memory_block="ASYNC SNAPSHOT FROM GATEWAY", user_block=None)

    async def archive_run_async(self, **kwargs) -> dict:
        self.archive_calls.append(kwargs)
        if self.archive_started is not None:
            self.archive_started.set()
        if self.archive_release is not None:
            await self.archive_release.wait()
        return {"status": "success"}


def _bundle(provider: _RecordingProvider) -> ProviderBundle:
    return ProviderBundle(
        main_runtime=SimpleNamespace(model="stub-model", provider_name="stub"),
        main_provider=provider,
    )


@pytest.mark.asyncio
async def test_agent_loop_uses_async_memory_snapshot(tmp_path) -> None:
    memory_service = _FakeMemoryService(tmp_path / "memory" / "curated")
    provider = _RecordingProvider("final answer")
    loop = AgentLoop(loader=EngineLoader(workspace=tmp_path, memory_service=memory_service))

    result = await loop.process_direct(
        "current task",
        user_id="user-1",
        session_id="session-1",
        provider_bundle=_bundle(provider),
        include_skill_assembly=False,
        include_tools=False,
    )

    assert result.output_text == "final answer"
    assert memory_service.snapshot_calls == [
        {"user_id": "user-1", "session_id": "session-1", "query": "current task"}
    ]
    assert "ASYNC SNAPSHOT FROM GATEWAY" in provider.messages[0][0]["content"]
    loop.close()


@pytest.mark.asyncio
async def test_agent_loop_archives_memory_gateway_run_in_background(tmp_path) -> None:
    memory_service = _FakeMemoryService(tmp_path / "memory" / "curated")
    memory_service.archive_started = asyncio.Event()
    memory_service.archive_release = asyncio.Event()
    provider = _RecordingProvider("assistant final")
    loop = AgentLoop(loader=EngineLoader(workspace=tmp_path, memory_service=memory_service))

    result = await loop.process_direct(
        "user asks",
        user_id="user-1",
        session_id="session-1",
        provider_bundle=_bundle(provider),
        include_skill_assembly=False,
        include_tools=False,
    )

    assert result.finish_reason == "stop"
    await asyncio.wait_for(memory_service.archive_started.wait(), timeout=1)
    assert memory_service.archive_calls == [
        {
            "user_id": "user-1",
            "session_id": "session-1",
            "user_message": "user asks",
            "assistant_message": "assistant final",
        }
    ]

    memory_service.archive_release.set()
    await asyncio.sleep(0)
    events = loop.boot().session_manager.get_run_event_records(result.session_id, result.run_id)
    assert any(event.event_type == "memory_gateway_archive_completed" for event in events)
    loop.close()
