import datetime

from sqlalchemy import (
    BigInteger,
    Column,
    Integer,
    String,
)
from consts.consts import DEFAULT_MAX_EDA_REQUESTS_ORG_PER_DAY

from models.base import Base, BaseMixin, with_session


class RateLimit(Base, BaseMixin):
    __tablename__ = "rate_limit"

    identifier = Column(String(128), nullable=False, unique=True)
    count = Column(Integer, nullable=False)

    @classmethod
    def incr(cls, identifier):
        count = 0

        with with_session() as s:
            item = s.query(cls).filter(cls.identifier == identifier).with_for_update().first()
            if item:
                item.count += 1
            else:
                item = cls(identifier=identifier, count=1)
                s.add(item)

            count = item.count

        return count

    @classmethod
    def by_org_id_per_day(cls, org_id: int):
        now = datetime.datetime.now()
        date_str = now.strftime("%Y-%m-%d")

        identifier = f"{org_id}-{date_str}"

        return cls.incr(identifier)


class RateLimitWhitelist(Base, BaseMixin):
    __tablename__ = "rate_limit_whitelist"

    org_id = Column(String(64), nullable=False, unique=True)
    count = Column(Integer, nullable=False)

    @classmethod
    def create(cls, s, org_id, count):
        item = cls(org_id=org_id, count=count)
        s.add(item)
        s.flush()

        return item

    @classmethod
    def get_by_org_id(cls, s, org_id):
        return s.query(cls).filter(cls.org_id == org_id).first()

    @classmethod
    def get_max_count_by_org_id(cls, org_id):
        count = DEFAULT_MAX_EDA_REQUESTS_ORG_PER_DAY

        with with_session() as s:
            item = cls.get_by_org_id(s, org_id)
            if item:
                count = item.count
        return count