import pytest

from beaver.engine import EngineLoader
from beaver.foundation.config.schema import MemoryGatewayConfig
from beaver.memory.gateway.service import GatewayAugmentedMemoryService
from beaver.services.memory_service import MemoryService


class _FakeGatewayClient:
    async def ensure_user(self, user_id: str) -> str:
        return f"{user_id}-key"

    async def get_profile(self, **kwargs):
        return {
            "status": "success",
            "profile": {"summary": "用户喜欢拿铁。"},
            "items": [{"summary": "用户偏好中文回复。"}],
        }

    async def get_session_context(self, **kwargs):
        return {
            "status": "success",
            "context": {"latest_archive_overview": "上次讨论了 memory gateway 接入。"},
            "items": [{"summary": "用户要求保留本地 MEMORY.md。"}],
        }

    async def search(self, **kwargs):
        return {
            "status": "success",
            "items": [
                {
                    "source_backend": "openviking",
                    "text": "需要异步写入 /memory-system/messages。",
                }
            ],
        }


class _FailingGatewayClient(_FakeGatewayClient):
    async def ensure_user(self, user_id: str) -> str:
        raise RuntimeError("gateway unavailable")


@pytest.mark.asyncio
async def test_gateway_snapshot_keeps_local_memory_and_adds_gateway_sections(tmp_path) -> None:
    local = MemoryService(tmp_path / "memory" / "curated")
    local.initialize()
    local.get_store().add("memory", "本地项目约定：默认用中文解释。")
    local.get_store().add("user", "本地用户画像：用户关注记忆系统。")
    service = GatewayAugmentedMemoryService(
        local_service=local,
        client=_FakeGatewayClient(),
        config=MemoryGatewayConfig(snapshot_search_limit=3),
    )

    snapshot = await service.capture_snapshot_for_run_async(
        user_id="user-1",
        session_id="session-1",
        query="如何接入 memory gateway",
    )
    prompt = "\n".join(snapshot.as_prompt_sections())

    assert "本地项目约定：默认用中文解释。" in prompt
    assert "本地用户画像：用户关注记忆系统。" in prompt
    assert "GATEWAY USER PROFILE" in prompt
    assert "用户喜欢拿铁。" in prompt
    assert "GATEWAY SESSION CONTEXT" in prompt
    assert "上次讨论了 memory gateway 接入。" in prompt
    assert "GATEWAY SEARCH RESULTS" in prompt
    assert "需要异步写入 /memory-system/messages。" in prompt


@pytest.mark.asyncio
async def test_gateway_snapshot_falls_back_to_local_memory_on_gateway_failure(tmp_path) -> None:
    local = MemoryService(tmp_path / "memory" / "curated")
    local.initialize()
    local.get_store().add("memory", "本地记忆仍然可用。")
    service = GatewayAugmentedMemoryService(
        local_service=local,
        client=_FailingGatewayClient(),
        config=MemoryGatewayConfig(snapshot_search_limit=3),
    )

    snapshot = await service.capture_snapshot_for_run_async(
        user_id="user-1",
        session_id="session-1",
        query="任何问题",
    )
    prompt = "\n".join(snapshot.as_prompt_sections())

    assert "本地记忆仍然可用。" in prompt
    assert "GATEWAY USER PROFILE" not in prompt


def test_engine_loader_uses_gateway_memory_service_without_replacing_tools(tmp_path) -> None:
    config_path = tmp_path / "config.json"
    config_path.write_text(
        """
        {
          "agents": {"defaults": {"workspace": "%s"}},
          "memory": {
            "mode": "gateway",
            "gateway": {"baseUrl": "http://gateway.test", "defaultUserId": "default-user"}
          }
        }
        """
        % str(tmp_path / "workspace"),
        encoding="utf-8",
    )

    loader = EngineLoader(config_path=config_path)
    loaded = loader.load()

    assert isinstance(loaded.memory_service, GatewayAugmentedMemoryService)
    assert "memory" in loaded.tools
    assert "session_search" in loaded.tools
    loaded.close()
