from peewee import fn

from ..models import FSMData


class FSMDataRepository:
    def get(self, key: str):
        return FSMData.get_or_none(key=key)

    def set_state(self, key: str, state: str):
        FSMData.insert(
            key=key,
            state=state,
            data=fn.COALESCE(
                FSMData.select(FSMData.data).where(FSMData.key == key), None
            ),
        ).on_conflict(
            conflict_target=[FSMData.key],
            update={FSMData.state: state},
        ).execute()

    def set_data(self, key: str, data: str):
        FSMData.insert(
            key=key,
            data=data,
            state=fn.COALESCE(
                FSMData.select(FSMData.state).where(FSMData.key == key), None
            ),
        ).on_conflict(
            conflict_target=[FSMData.key],
            update={FSMData.data: data},
        ).execute()