from uuid import UUID

from sqlalchemy import and_, not_, select, update
from sqlalchemy.ext.asyncio import async_sessionmaker

from src.models.database_models.model import RefreshTokens, engine
from src.models.pydantic_models.model import RefreshTokensOut


class JwtCrudActions:
    def __init__(self) -> None:
        self.Session=async_sessionmaker(bind=engine)
    
    async def get_token_by_user_id(self, user_id:UUID)->RefreshTokensOut|None:
        async with self.Session() as session, session.begin():
            query=select(RefreshTokens).where(and_(RefreshTokens.user_id==user_id, not_(RefreshTokens.is_revoked)))
            response= (await session.scalars(query)).first()
            if response is None:
                return None
            return RefreshTokensOut.model_validate(response)
            
    async def get_token_by_id(self, token_id:UUID)->RefreshTokensOut|None:
        async with self.Session() as session, session.begin():
            query=select(RefreshTokens).where(RefreshTokens.id==token_id)
            response= (await session.scalars(query)).one_or_none()
            if response is None:
                return None
            return RefreshTokensOut.model_validate(response)
    
    async def create_token(self, data:dict)->None:
        async with self.Session() as session, session.begin():
            new_token=RefreshTokens(**data)
            session.add(new_token)
        
            
    async def create_and_update_token(self, data: dict, old_jti: UUID, new_jti: UUID) -> bool:
        async with self.Session() as session, session.begin():
            new_token = RefreshTokens(**data)

            query = (
                update(RefreshTokens)
                .where(RefreshTokens.id == old_jti, RefreshTokens.is_revoked.is_(False))
                .values(is_revoked=True, replaced_by=new_jti)
                .returning(RefreshTokens.id)
            )
            updated_id = (await session.execute(query)).scalar_one_or_none()

            if updated_id is None:
                return False

            session.add(new_token)
            return True
                
    
    async def revoke_all(self, user_id:UUID)->bool:
        async with self.Session() as session, session.begin():
            await session.execute(update(RefreshTokens).where(RefreshTokens.user_id==user_id).values(is_revoked=True)) #bulk update
            return True

    async def logout(self,token_id:UUID)->bool:
        async with self.Session() as session, session.begin():
            query=select(RefreshTokens).where(RefreshTokens.id == token_id)
            response= (await session.scalars(query)).one_or_none()
            if response is None:
                return False
            else:
                response.is_revoked=True
                return True