from __future__ import annotations

from datetime import datetime

from sqlalchemy import delete, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.enums import ThrottleAction
from app.models.ip_ban import IpBan
from app.models.ip_throttle_state import IpThrottleState


class IpThrottleRepository:
    def __init__(self, db: AsyncSession) -> None:
        self.db = db

    async def get_state(self, ip_address: str, action: ThrottleAction) -> IpThrottleState | None:
        result = await self.db.execute(
            select(IpThrottleState).where(
                IpThrottleState.ip_address == ip_address, IpThrottleState.action == action
            )
        )
        return result.scalar_one_or_none()

    async def get_or_create_state(self, ip_address: str, action: ThrottleAction) -> IpThrottleState:
        state = await self.get_state(ip_address, action)
        if state is not None:
            return state
        state = IpThrottleState(ip_address=ip_address, action=action)
        self.db.add(state)
        await self.db.flush()
        return state

    async def get_ban(self, ip_address: str) -> IpBan | None:
        result = await self.db.execute(select(IpBan).where(IpBan.ip_address == ip_address))
        return result.scalar_one_or_none()

    async def create_ban(self, ip_address: str, reason: str, banned_at: datetime) -> IpBan:
        ban = IpBan(ip_address=ip_address, reason=reason, banned_at=banned_at)
        self.db.add(ban)
        await self.db.flush()
        return ban

    async def list_bans(self) -> list[IpBan]:
        result = await self.db.execute(select(IpBan).order_by(IpBan.banned_at.desc()))
        return list(result.scalars().all())

    async def clear_ban_and_state(self, ip_address: str) -> bool:
        """Full pardon for an admin-approved unban - removes the ban record
        and resets every throttle-state row for this IP to a clean slate
        (not just lifting the terminal ban while leaving them one offense
        from another)."""
        ban_result = await self.db.execute(delete(IpBan).where(IpBan.ip_address == ip_address))
        result = await self.db.execute(
            select(IpThrottleState).where(IpThrottleState.ip_address == ip_address)
        )
        states = result.scalars().all()
        for state in states:
            state.attempt_count = 0
            state.next_allowed_at = None
            state.timeout_until = None
            state.offense_count = 0
        await self.db.flush()
        return ban_result.rowcount > 0 or len(states) > 0