from datetime import datetime, timezone

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.db.models import AssistantState


def get_state(db: Session, key: str) -> str | None:
    row = db.get(AssistantState, key)
    return row.value if row else None


def set_state(db: Session, key: str, value: str) -> None:
    row = db.get(AssistantState, key)
    now = datetime.now(timezone.utc)
    if row:
        row.value = value
        row.updated_at = now
    else:
        db.add(AssistantState(key=key, value=value, updated_at=now))
    db.commit()