from uuid import UUID

from sqlalchemy import select

from app.models.membership import TenantMembership
from app.models.user import UserProfile


class UserRepository:
    def __init__(self, session) -> None:
        self.session = session

    async def get(self, user_id: UUID) -> UserProfile | None:
        return await self.session.get(UserProfile, user_id)

    async def get_by_sub(self, keycloak_sub: str) -> UserProfile | None:
        stmt = select(UserProfile).where(UserProfile.keycloak_sub == keycloak_sub)
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def get_by_mobile(self, mobile: str) -> UserProfile | None:
        stmt = select(UserProfile).where(UserProfile.mobile == mobile)
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def get_by_username(self, username: str) -> UserProfile | None:
        stmt = select(UserProfile).where(UserProfile.username == username)
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def add(self, user: UserProfile) -> UserProfile:
        self.session.add(user)
        await self.session.flush()
        return user

    async def list(self, offset: int = 0, limit: int = 20) -> list[UserProfile]:
        stmt = (
            select(UserProfile)
            .order_by(UserProfile.created_at.desc())
            .offset(offset)
            .limit(limit)
        )
        result = await self.session.execute(stmt)
        return list(result.scalars().all())


class MembershipRepository:
    def __init__(self, session) -> None:
        self.session = session

    async def get(self, tenant_id: UUID, user_id: UUID) -> TenantMembership | None:
        stmt = select(TenantMembership).where(
            TenantMembership.tenant_id == tenant_id,
            TenantMembership.user_id == user_id,
        )
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def list_by_tenant(self, tenant_id: UUID) -> list[TenantMembership]:
        stmt = select(TenantMembership).where(TenantMembership.tenant_id == tenant_id)
        result = await self.session.execute(stmt)
        return list(result.scalars().all())

    async def list_by_user(self, user_id: UUID) -> list[TenantMembership]:
        stmt = select(TenantMembership).where(TenantMembership.user_id == user_id)
        result = await self.session.execute(stmt)
        return list(result.scalars().all())

    async def add(self, membership: TenantMembership) -> TenantMembership:
        self.session.add(membership)
        await self.session.flush()
        return membership