from sqlalchemy import select
from sqlalchemy.orm import Session

from app.models.entities import InviteCode, InviteRelation, User


class InviteRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def get_default_code(self, user_id: int) -> InviteCode | None:
        return self.db.scalar(
            select(InviteCode).where(
                InviteCode.user_id == user_id,
                InviteCode.is_default.is_(True),
            )
        )

    def get_code(self, code_value: str) -> InviteCode | None:
        return self.db.scalar(select(InviteCode).where(InviteCode.invite_code == code_value))

    def inviter_relations(self, user_id: int) -> list[InviteRelation]:
        return (
            self.db.query(InviteRelation)
            .filter(InviteRelation.inviter_user_id == user_id)
            .order_by(InviteRelation.id.desc())
            .all()
        )

    def users_by_ids(self, user_ids: list[int]) -> dict[int, User]:
        if not user_ids:
            return {}
        rows = self.db.scalars(select(User).where(User.id.in_(user_ids))).all()
        return {row.id: row for row in rows}