from contextlib import asynccontextmanager

from aiogram.fsm.context import FSMContext
from pydantic.main import BaseModel
from typing_extensions import AsyncGenerator


async def set_data(state: FSMContext, model: BaseModel) -> None:
    await state.set_data(model.model_dump())


async def get_data[T: BaseModel](state: FSMContext, model_type: type[T]) -> T:
    return model_type.model_validate(await state.get_data())


@asynccontextmanager
async def edit_data[T: BaseModel](
    state: FSMContext,
    model_type: type[T],
) -> AsyncGenerator[T]:
    model = await get_data(state, model_type)
    yield model
    await set_data(state, model)