social-app/backend/src/v1/friendships/repository.py

from __future__ import annotations

from datetime import datetime, timezone
from typing import TYPE_CHECKING, Protocol
from uuid import UUID

from sqlalchemy import select, or_
from sqlalchemy.exc import SQLAlchemyError

from core.db.base_repository import BaseRepository
from core.logging import get_logger
from models.friendships import Friendship, FriendshipStatus
from models.inbox_messages import InboxMessage, InboxMessageStatus, InboxMessageType
from schemas.inbox.messages import FriendshipContent

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

logger = get_logger("v1.friendships.repository")


class FriendshipRepository(Protocol):
    """Protocol defining the friendship repository interface."""

    async def create_request(
        self, initiator_id: UUID, recipient_id: UUID, content: str | None = None
    ) -> tuple[Friendship, InboxMessage]:
        """Create a friendship request and inbox message."""
        ...

    async def reactivate_request(
        self,
        friendship: Friendship,
        initiator_id: UUID,
        content: str | None = None,
    ) -> tuple[Friendship, InboxMessage]:
        """Reactivate a declined or canceled friendship request."""
        ...

    async def get_friendship_between_users(
        self, user_id_1: UUID, user_id_2: UUID
    ) -> Friendship | None:
        """Check if friendship exists between two users."""
        ...

    async def get_pending_inbox_for_recipient(
        self, recipient_id: UUID, friendship_id: UUID
    ) -> InboxMessage | None:
        """Get pending inbox message for recipient."""
        ...

    async def get_friendship_by_id(self, friendship_id: UUID) -> Friendship | None:
        """Get friendship by ID."""
        ...

    async def get_friendships_by_ids(
        self, friendship_ids: list[UUID]
    ) -> dict[UUID, Friendship]:
        """Batch get friendships by IDs."""
        ...

    async def get_inbox_messages_for_user(
        self, user_id: UUID, status: InboxMessageStatus | None = None
    ) -> list[InboxMessage]:
        """Get inbox messages for a user."""
        ...

    async def get_outgoing_requests(self, user_id: UUID) -> list[Friendship]:
        """Get outgoing friend requests."""
        ...

    async def get_friends_list(self, user_id: UUID) -> list[Friendship]:
        """Get accepted friends list."""
        ...


class SQLAlchemyFriendshipRepository(BaseRepository[Friendship]):
    """SQLAlchemy implementation of FriendshipRepository.

    Note: This repository only performs CRUD operations.
    - No commit (only flush) - service layer handles transactions
    - No auth logic - service layer handles authorization
    - No HTTP exceptions - returns None or raises SQLAlchemyError
    """

    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session, Friendship)

    async def create_request(
        self, initiator_id: UUID, recipient_id: UUID, content: str | None = None
    ) -> tuple[Friendship, InboxMessage]:
        try:
            user_low_id = min(initiator_id, recipient_id)
            user_high_id = max(initiator_id, recipient_id)
            now = datetime.now(timezone.utc)

            friendship = Friendship(
                user_low_id=user_low_id,
                user_high_id=user_high_id,
                initiator_id=initiator_id,
                status=FriendshipStatus.PENDING,
                requested_at=now,
                created_by=initiator_id,
                updated_by=initiator_id,
            )
            self._session.add(friendship)
            await self._session.flush()

            inbox_content = FriendshipContent(type="request", message=content)
            inbox = InboxMessage(
                recipient_id=recipient_id,
                sender_id=initiator_id,
                message_type=InboxMessageType.FRIEND_REQUEST,
                friendship_id=friendship.id,
                content=inbox_content.model_dump(),
                status=InboxMessageStatus.PENDING,
                created_by=initiator_id,
            )
            self._session.add(inbox)
            await self._session.flush()

            return friendship, inbox
        except SQLAlchemyError:
            logger.exception(
                "Failed to create friendship request",
                initiator_id=str(initiator_id),
                recipient_id=str(recipient_id),
            )
            raise

    async def reactivate_request(
        self,
        friendship: Friendship,
        initiator_id: UUID,
        content: str | None = None,
    ) -> tuple[Friendship, InboxMessage]:
        try:
            now = datetime.now(timezone.utc)
            friendship.status = FriendshipStatus.PENDING
            friendship.requested_at = now
            friendship.initiator_id = initiator_id
            friendship.updated_by = initiator_id

            inbox_content = FriendshipContent(type="request", message=content)
            inbox = InboxMessage(
                recipient_id=(
                    friendship.user_low_id
                    if initiator_id == friendship.user_high_id
                    else friendship.user_high_id
                ),
                sender_id=initiator_id,
                message_type=InboxMessageType.FRIEND_REQUEST,
                friendship_id=friendship.id,
                content=inbox_content.model_dump(),
                status=InboxMessageStatus.PENDING,
                created_by=initiator_id,
            )
            self._session.add(inbox)
            await self._session.flush()

            return friendship, inbox
        except SQLAlchemyError:
            logger.exception(
                "Failed to reactivate friendship request",
                friendship_id=str(friendship.id),
                initiator_id=str(initiator_id),
            )
            raise

    async def get_friendship_between_users(
        self, user_id_1: UUID, user_id_2: UUID
    ) -> Friendship | None:
        try:
            user_low_id = min(user_id_1, user_id_2)
            user_high_id = max(user_id_1, user_id_2)

            stmt = (
                select(Friendship)
                .where(Friendship.user_low_id == user_low_id)
                .where(Friendship.user_high_id == user_high_id)
                .where(Friendship.deleted_at.is_(None))
            )
            result = await self._session.execute(stmt)
            return result.scalar_one_or_none()
        except SQLAlchemyError:
            logger.exception(
                "Failed to get friendship between users",
                user_id_1=str(user_id_1),
                user_id_2=str(user_id_2),
            )
            raise

    async def get_pending_inbox_for_recipient(
        self, recipient_id: UUID, friendship_id: UUID
    ) -> InboxMessage | None:
        try:
            stmt = (
                select(InboxMessage)
                .where(InboxMessage.recipient_id == recipient_id)
                .where(InboxMessage.friendship_id == friendship_id)
                .where(InboxMessage.status == InboxMessageStatus.PENDING)
            )
            result = await self._session.execute(stmt)
            return result.scalar_one_or_none()
        except SQLAlchemyError:
            logger.exception(
                "Failed to get pending inbox message",
                recipient_id=str(recipient_id),
                friendship_id=str(friendship_id),
            )
            raise

    async def get_friendship_by_id(self, friendship_id: UUID) -> Friendship | None:
        try:
            return await self.get_by_id(friendship_id)
        except SQLAlchemyError:
            logger.exception(
                "Failed to get friendship by id",
                friendship_id=str(friendship_id),
            )
            raise

    async def get_friendships_by_ids(
        self, friendship_ids: list[UUID]
    ) -> dict[UUID, Friendship]:
        if not friendship_ids:
            return {}
        try:
            unique_ids = list(dict.fromkeys(friendship_ids))
            stmt = (
                select(Friendship)
                .where(Friendship.id.in_(unique_ids))
                .where(Friendship.deleted_at.is_(None))
            )
            result = await self._session.execute(stmt)
            friendships = list(result.scalars().all())
            return {friendship.id: friendship for friendship in friendships}
        except SQLAlchemyError:
            logger.exception(
                "Failed to get friendships by ids",
                friendship_ids=[str(i) for i in friendship_ids],
            )
            raise

    async def get_inbox_messages_for_user(
        self, user_id: UUID, status: InboxMessageStatus | None = None
    ) -> list[InboxMessage]:
        try:
            stmt = (
                select(InboxMessage)
                .where(InboxMessage.recipient_id == user_id)
                .order_by(InboxMessage.created_at.desc())
            )
            if status is not None:
                stmt = stmt.where(InboxMessage.status == status)

            result = await self._session.execute(stmt)
            return list(result.scalars().all())
        except SQLAlchemyError:
            logger.exception(
                "Failed to get inbox messages for user",
                user_id=str(user_id),
            )
            raise

    async def get_outgoing_requests(self, user_id: UUID) -> list[Friendship]:
        try:
            stmt = (
                select(Friendship)
                .where(Friendship.initiator_id == user_id)
                .where(Friendship.status == FriendshipStatus.PENDING)
                .where(Friendship.deleted_at.is_(None))
                .order_by(Friendship.created_at.desc())
            )
            result = await self._session.execute(stmt)
            return list(result.scalars().all())
        except SQLAlchemyError:
            logger.exception(
                "Failed to get outgoing requests",
                user_id=str(user_id),
            )
            raise

    async def get_friends_list(self, user_id: UUID) -> list[Friendship]:
        try:
            stmt = (
                select(Friendship)
                .where(
                    or_(
                        Friendship.user_low_id == user_id,
                        Friendship.user_high_id == user_id,
                    )
                )
                .where(Friendship.status == FriendshipStatus.ACCEPTED)
                .where(Friendship.deleted_at.is_(None))
                .order_by(Friendship.updated_at.desc())
            )
            result = await self._session.execute(stmt)
            return list(result.scalars().all())
        except SQLAlchemyError:
            logger.exception(
                "Failed to get friends list",
                user_id=str(user_id),
            )
            raise