"""``SessionStore`` adapter for the Claude Agent SDK on top of :class:`Database`.

Entries are stored verbatim as JSON (``jsonb`` on Postgres), ordered by a
per-key ``seq``. ``append`` is idempotent on ``entry["uuid"]`` because the
SDK re-delivers a batch on retry; entries without a uuid are appended as
they come. Runs on SQLite too - the conformance suite uses that in tests.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, cast

from claude_agent_sdk import SessionStore
from sqlalchemy import delete, func
from sqlmodel import col, select

from beaver_gateway.storage.models import TranscriptEntry

if TYPE_CHECKING:
    from claude_agent_sdk.types import (
        SessionKey,
        SessionListSubkeysKey,
        SessionStoreEntry,
    )

    from beaver_gateway.storage.db import Database

__all__ = ["PostgresSessionStore"]


class PostgresSessionStore(SessionStore):
    def __init__(self, db: Database) -> None:
        self._db = db

    async def append(self, key: SessionKey, entries: list[SessionStoreEntry]) -> None:
        if not entries:
            return
        project_key, session_id, subpath = _parts(key)
        wanted = [
            u for u in (e.get("uuid") for e in entries) if isinstance(u, str) and u
        ]
        async with self._db.session() as session:
            known: set[str] = set()
            if wanted:
                result = await session.exec(
                    select(TranscriptEntry.uuid).where(
                        TranscriptEntry.project_key == project_key,
                        TranscriptEntry.session_id == session_id,
                        TranscriptEntry.subpath == subpath,
                        col(TranscriptEntry.uuid).in_(wanted),
                    )
                )
                known = {u for u in result.all() if u is not None}
            seq_result = await session.exec(
                select(func.max(TranscriptEntry.seq)).where(
                    TranscriptEntry.project_key == project_key,
                    TranscriptEntry.session_id == session_id,
                    TranscriptEntry.subpath == subpath,
                )
            )
            seq = seq_result.one() or 0
            for entry in entries:
                uuid = entry.get("uuid")
                uuid_str = uuid if isinstance(uuid, str) and uuid else None
                if uuid_str is not None:
                    if uuid_str in known:
                        continue
                    known.add(uuid_str)
                seq += 1
                session.add(
                    TranscriptEntry(
                        project_key=project_key,
                        session_id=session_id,
                        subpath=subpath,
                        seq=seq,
                        uuid=uuid_str,
                        entry=cast("dict[str, Any]", dict(entry)),
                    )
                )
            await session.commit()

    async def load(self, key: SessionKey) -> list[SessionStoreEntry] | None:
        project_key, session_id, subpath = _parts(key)
        async with self._db.session() as session:
            result = await session.exec(
                select(TranscriptEntry.entry)
                .where(
                    TranscriptEntry.project_key == project_key,
                    TranscriptEntry.session_id == session_id,
                    TranscriptEntry.subpath == subpath,
                )
                .order_by(col(TranscriptEntry.seq))
            )
            rows = result.all()
        if not rows:
            return None
        return [cast("SessionStoreEntry", row) for row in rows]

    async def list_subkeys(self, key: SessionListSubkeysKey) -> list[str]:
        async with self._db.session() as session:
            result = await session.exec(
                select(TranscriptEntry.subpath)
                .where(
                    TranscriptEntry.project_key == key["project_key"],
                    TranscriptEntry.session_id == key["session_id"],
                    TranscriptEntry.subpath != "",
                )
                .distinct()
            )
            return sorted(result.all())

    async def delete(self, key: SessionKey) -> None:
        project_key, session_id, subpath = _parts(key)
        stmt = delete(TranscriptEntry).where(
            col(TranscriptEntry.project_key) == project_key,
            col(TranscriptEntry.session_id) == session_id,
        )
        if key.get("subpath"):
            stmt = stmt.where(col(TranscriptEntry.subpath) == subpath)
        async with self._db.session() as session:
            await session.execute(stmt)  # ty: ignore[deprecated]
            await session.commit()


def _parts(key: SessionKey) -> tuple[str, str, str]:
    return key["project_key"], key["session_id"], key.get("subpath") or ""