"""Restore a conversation through a new SQLite connection and a new Agent."""

import asyncio
from collections.abc import AsyncIterator
from pathlib import Path
from tempfile import TemporaryDirectory

from zett_agent.agent import AgentRunConfig
from zett_agent.client import create_agent
from zett_agent.extensions.sqlite import SQLiteSessionExtension
from zett_agent.messages import (
    AssistantMessage,
    UserMessage,
)
from zett_agent.model import (
    ModelEvent,
    ModelRequest,
    ModelResponse,
    RetryOptions,
)


class TurnModel:
    retry = RetryOptions(max_retries=0)

    async def stream(self, request: ModelRequest) -> AsyncIterator[ModelEvent]:
        turn = sum(isinstance(message, UserMessage) for message in request.messages)
        yield ModelEvent.completed(ModelResponse(AssistantMessage(content=f"Stored turn {turn}")))


async def main() -> None:
    # No demo touches the default database under the user's home directory.
    with TemporaryDirectory(prefix="zett-docs-") as directory:
        path = Path(directory) / "sessions.sqlite"
        session_id = "project-notes"
        config = AgentRunConfig(session_id=session_id)
        first_storage = SQLiteSessionExtension(path)
        try:
            first = await create_agent(TurnModel(), config=config, extensions=[first_storage])
            print((await first.run("Remember Python", metadata={"origin": "tutorial"})).content)
        finally:
            await first_storage.close()

        restored_storage = SQLiteSessionExtension(path)
        try:
            restored = await create_agent(TurnModel(), config=config, extensions=[restored_storage])
            print((await restored.run("Continue")).content)
            rows = await restored_storage.list_raw_messages(session_id)
            assert [row.message.role.value for row in rows] == [
                "system",
                "user",
                "assistant",
                "system",
                "user",
                "assistant",
            ]
            assert rows[0].metadata == {"origin": "tutorial"}
            assert len(await restored_storage.list_raw_messages(session_id, limit=2, offset=2)) == 2
            view = await restored_storage.storage.load(session_id)
            assert view.snapshot is None and len(view.raw_tail) == 6
            await restored_storage.update_session(session_id, title="Project notes")
            renamed = await restored_storage.get_session(session_id)
            assert renamed is not None and renamed.title == "Project notes"
            assert await restored_storage.delete_session(session_id)
            assert await restored_storage.list_raw_messages(session_id) == []
            print("Raw roles: system, user, assistant, system, user, assistant")
            print("Session restored, renamed, and deleted in a temporary database")
        finally:
            await restored_storage.close()


if __name__ == "__main__":
    asyncio.run(main())