"""Integration tests for agent memory prompt preparation."""

from collections.abc import AsyncGenerator
from typing import Any, Optional

import pytest

from agentbyte.agents import Agent
from agentbyte.llm.base import BaseChatCompletionClient, BaseChatCompletionClientConfig
from agentbyte.llm.types import ChatCompletionChunk, ChatCompletionResult
from agentbyte.memory import ListMemory, MemoryContent
from agentbyte.messages import AssistantMessage, Message, Usage, UserMessage


class FakeClientConfig(BaseChatCompletionClientConfig):
    pass


class CaptureMessagesClient(BaseChatCompletionClient):
    def __init__(self) -> None:
        super().__init__(model="fake", client=object())
        self.last_messages: list[Message] = []

    async def create(
        self,
        messages: list[Message],
        tools: Optional[list[dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        self.last_messages = messages
        return ChatCompletionResult(
            message=AssistantMessage(content="done", source="fake"),
            usage=Usage(tokens_input=1, tokens_output=1),
            model="fake",
        )

    async def create_stream(
        self,
        messages: list[Message],
        tools: Optional[list[dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        if False:
            yield ChatCompletionChunk(content="", model="fake", is_complete=True)

    def _to_config(self) -> BaseChatCompletionClientConfig:
        return FakeClientConfig(model=self.model, config=self.config)

    @classmethod
    def _from_config(
        cls,
        config: BaseChatCompletionClientConfig,
    ) -> "CaptureMessagesClient":
        return cls()


@pytest.mark.asyncio
async def test_agent_injects_memory_into_system_prompt() -> None:
    client = CaptureMessagesClient()
    memory = ListMemory()
    await memory.add(MemoryContent(content="User prefers concise responses"))

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="You are helpful.",
        model_client=client,
        memory=memory,
    )

    await agent.run(UserMessage(content="hello", source="user"))

    assert client.last_messages
    system_message = client.last_messages[0]
    assert "Relevant context from memory" in system_message.content
    assert "User prefers concise responses" in system_message.content


@pytest.mark.asyncio
async def test_agent_without_memory_keeps_base_system_prompt() -> None:
    client = CaptureMessagesClient()

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="Base instructions.",
        model_client=client,
    )

    await agent.run(UserMessage(content="hello", source="user"))

    assert client.last_messages
    system_message = client.last_messages[0]
    assert system_message.content == "Base instructions."
