"""aiogram FSM storage on top of the `fsm_state` table — dialog state survives bot restarts (SPEC M-4)."""

from collections.abc import Mapping
from typing import Any

from aiogram.fsm.state import State
from aiogram.fsm.storage.base import BaseStorage, StateType, StorageKey

from app.db.models import FsmState
from app.db.session import SessionMaker


def key_to_str(key: StorageKey) -> str:
    parts = (key.bot_id, key.chat_id, key.user_id, key.thread_id, key.business_connection_id, key.destiny)
    return ":".join("" if p is None else str(p) for p in parts)


class SqlStorage(BaseStorage):
    def __init__(self, sessions: SessionMaker) -> None:
        self._sessions = sessions

    async def set_state(self, key: StorageKey, state: StateType = None) -> None:
        value = state.state if isinstance(state, State) else state
        await self._write(key, state=value)

    async def get_state(self, key: StorageKey) -> str | None:
        async with self._sessions() as session:
            row = await session.get(FsmState, key_to_str(key))
            return row.state if row else None

    async def set_data(self, key: StorageKey, data: Mapping[str, Any]) -> None:
        await self._write(key, data=dict(data))

    async def get_data(self, key: StorageKey) -> dict[str, Any]:
        async with self._sessions() as session:
            row = await session.get(FsmState, key_to_str(key))
            return dict(row.data) if row else {}

    async def close(self) -> None:
        pass  # the engine is owned and disposed by the app

    async def _write(self, key: StorageKey, **fields: Any) -> None:
        async with self._sessions.begin() as session:
            row = await session.get(FsmState, key_to_str(key))
            state = fields.get("state", row.state if row else None)
            data = fields.get("data", row.data if row else {})
            if state is None and not data:
                if row is not None:
                    await session.delete(row)  # empty state: keep the table small
                return
            if row is None:
                session.add(FsmState(key=key_to_str(key), state=state, data=data))
            else:
                row.state, row.data = state, data