from typing import Union
from sqlalchemy import (
    Column,
    Integer,
    String,
)

from models.base import Base, BaseMixin, with_session


class OpenAIApiKey(Base, BaseMixin):
    __tablename__ = 'openai_api_keys'

    app_id = Column(Integer, nullable=False, unique=True)
    openai_api_key = Column(String(128), nullable=False)

    @classmethod
    def upsert(cls, s, app_id, openai_api_key):
        key = cls(app_id=app_id, openai_api_key=openai_api_key)
        if not key:
            s.add(key)
        else:
            key.openai_api_key = openai_api_key
        s.flush()
        return key

    @classmethod
    def get_api_key_by_app_id(cls, app_id: int) -> Union[str, None]:
        with with_session() as s:
            item = s.query(cls).filter(cls.app_id == app_id).first()
            if not item:
                return None

            return item.openai_api_key

    @classmethod
    def delete_by_app_id(cls, s, app_id: int) -> None:
        item = s.query(cls).filter(cls.app_id == app_id).first()
        if not item:
            return

        s.delete(item)
        s.flush()