from __future__ import annotations

from typing import TYPE_CHECKING, Any

from sqlalchemy import (  # Dishka needs runtime access
    Table,
    func as sa_func,
    select,
)
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession  # noqa: TC002  # Dishka needs runtime access

from waku.eventsourcing.contracts.stream import StreamId
from waku.eventsourcing.snapshot.interfaces import ISnapshotStore, Snapshot

if TYPE_CHECKING:
    from collections.abc import Callable

__all__ = [
    'SqlAlchemySnapshotStore',
    'make_sqlalchemy_snapshot_store',
]


class SqlAlchemySnapshotStore(ISnapshotStore):
    def __init__(self, session: AsyncSession, snapshots_table: Table) -> None:
        self._session = session
        self._snapshots = snapshots_table

    async def load(self, stream_id: StreamId, /) -> Snapshot | None:
        key = str(stream_id)
        query = select(self._snapshots).where(self._snapshots.c.stream_id == key)
        result = await self._session.execute(query)
        row: Any = result.one_or_none()
        if row is None:
            return None
        return Snapshot(
            stream_id=StreamId.from_value(row.stream_id),
            state=row.state,
            version=row.version,
            state_type=row.state_type,
            schema_version=row.schema_version,
        )

    async def save(self, snapshot: Snapshot, /) -> None:
        stmt = pg_insert(self._snapshots).values(
            stream_id=str(snapshot.stream_id),
            state=snapshot.state,
            version=snapshot.version,
            state_type=snapshot.state_type,
            schema_version=snapshot.schema_version,
        )
        stmt = stmt.on_conflict_do_update(
            index_elements=['stream_id'],
            set_={
                'state': stmt.excluded.state,
                'version': stmt.excluded.version,
                'state_type': stmt.excluded.state_type,
                'schema_version': stmt.excluded.schema_version,
                'updated_at': sa_func.now(),
            },
        )
        await self._session.execute(stmt)
        await self._session.flush()


def make_sqlalchemy_snapshot_store(
    snapshots_table: Table,
) -> Callable[..., SqlAlchemySnapshotStore]:
    def factory(session: AsyncSession) -> SqlAlchemySnapshotStore:
        return SqlAlchemySnapshotStore(session, snapshots_table)

    return factory
