"""Message persistence for conversation history."""

import json
from typing import Any

from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter


class MessageStore:
    """Store and retrieve conversation messages.

    This class handles serialization/deserialization of PydanticAI messages
    for persistence in a database or other storage backend.

    Example:
        ```python
        store = MessageStore(database_pool)

        # Save messages after agent run
        await store.save_messages(
            conversation_id="conv-123",
            messages=result.new_messages()
        )

        # Load messages for next run
        messages = await store.load_messages("conv-123")
        result = await agent.run("query", message_history=messages)
        ```
    """

    def __init__(self, db: Any) -> None:
        """Initialize the message store.

        Args:
            db: Database connection pool with fetch_all and execute methods.
        """
        self.db = db

    async def save_messages(
        self,
        conversation_id: str,
        messages: list[ModelMessage],
        user_id: str | None = None,
    ) -> None:
        """Save messages to the database.

        Args:
            conversation_id: Unique identifier for the conversation.
            messages: List of messages to save.
            user_id: The user's ID (for listing conversations).
        """
        # Serialize messages using PydanticAI's type adapter
        messages_json = ModelMessagesTypeAdapter.dump_json(messages).decode()

        await self.db.execute(
            """
            INSERT INTO conversation_messages (conversation_id, user_id, messages, created_at)
            VALUES ($1, $2, $3, NOW())
            ON CONFLICT (conversation_id)
            DO UPDATE SET messages = conversation_messages.messages || $3::jsonb,
                          updated_at = NOW()
            """,
            conversation_id,
            user_id,
            messages_json,
        )

    async def load_messages(
        self,
        conversation_id: str,
    ) -> list[ModelMessage]:
        """Load messages from the database.

        Args:
            conversation_id: Unique identifier for the conversation.

        Returns:
            List of messages, or empty list if conversation not found.
        """
        result = await self.db.fetch_one(
            "SELECT messages FROM conversation_messages WHERE conversation_id = $1",
            conversation_id,
        )

        if not result:
            return []

        # Deserialize using PydanticAI's type adapter
        messages_data = result["messages"]
        if isinstance(messages_data, str):
            messages_data = json.loads(messages_data)

        return ModelMessagesTypeAdapter.validate_python(messages_data)

    async def delete_messages(
        self,
        conversation_id: str,
    ) -> bool:
        """Delete all messages for a conversation.

        Args:
            conversation_id: Unique identifier for the conversation.

        Returns:
            True if messages were deleted, False if not found.
        """
        result = await self.db.fetch_one(
            """DELETE FROM conversation_messages
            WHERE conversation_id = $1 RETURNING conversation_id""",
            conversation_id,
        )
        return result is not None

    async def get_recent_conversations(
        self,
        user_id: str,
        limit: int = 20,
    ) -> list[dict[str, Any]]:
        """Get recent conversations for a user.

        Args:
            user_id: The user's ID.
            limit: Maximum number of conversations to return.

        Returns:
            List of conversation metadata (id, created_at, message_count).
        """
        rows = await self.db.fetch_all(
            """
            SELECT
                conversation_id,
                created_at,
                updated_at,
                jsonb_array_length(messages) as message_count
            FROM conversation_messages
            WHERE user_id = $1
            ORDER BY updated_at DESC
            LIMIT $2
            """,
            user_id,
            limit,
        )

        return [dict(row) for row in rows]
