mitlist/be/app/crud/invite.py

# app/crud/invite.py
import secrets
from datetime import datetime, timedelta, timezone
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy.orm import selectinload # Ensure selectinload is imported
from sqlalchemy import delete # Import delete statement
from sqlalchemy.exc import SQLAlchemyError, OperationalError, IntegrityError
from typing import Optional

from app.models import Invite as InviteModel, Group as GroupModel, User as UserModel # Import related models for selectinload
from app.core.exceptions import (
    DatabaseConnectionError,
    DatabaseIntegrityError,
    DatabaseQueryError,
    DatabaseTransactionError,
    InviteOperationError # Add new specific exception
)

# Invite codes should be reasonably unique, but handle potential collision
MAX_CODE_GENERATION_ATTEMPTS = 5

async def deactivate_all_active_invites_for_group(db: AsyncSession, group_id: int):
    """Deactivates all currently active invite codes for a specific group."""
    try:
        stmt = (
            select(InviteModel)
            .where(InviteModel.group_id == group_id, InviteModel.is_active == True)
        )
        result = await db.execute(stmt)
        active_invites = result.scalars().all()

        if not active_invites:
            return # No active invites to deactivate

        for invite in active_invites:
            invite.is_active = False
            db.add(invite)
        
        # await db.flush() # Removed: Rely on caller to flush/commit
        # No explicit commit here, assuming it's part of a larger transaction or caller handles commit.
    except OperationalError as e:
        # It's better to let the caller handle rollback or commit based on overall operation success
        raise DatabaseConnectionError(f"DB connection error deactivating invites for group {group_id}: {str(e)}")
    except SQLAlchemyError as e:
        raise DatabaseTransactionError(f"DB transaction error deactivating invites for group {group_id}: {str(e)}")

async def create_invite(db: AsyncSession, group_id: int, creator_id: int, expires_in_days: int = 365 * 100) -> Optional[InviteModel]: # Default to 100 years
    """Creates a new invite code for a group, deactivating any existing active ones for that group first."""
    
    # Deactivate existing active invites for this group
    await deactivate_all_active_invites_for_group(db, group_id)

    expires_at = datetime.now(timezone.utc) + timedelta(days=expires_in_days)
    
    potential_code = None
    for attempt in range(MAX_CODE_GENERATION_ATTEMPTS):
        potential_code = secrets.token_urlsafe(16)
        existing_check_stmt = select(InviteModel.id).where(InviteModel.code == potential_code, InviteModel.is_active == True).limit(1)
        existing_result = await db.execute(existing_check_stmt)
        if existing_result.scalar_one_or_none() is None:
            break
        if attempt == MAX_CODE_GENERATION_ATTEMPTS - 1:
            raise InviteOperationError("Failed to generate a unique invite code after several attempts.")

    # Removed explicit transaction block here, rely on FastAPI's request-level transaction.
    # Final check for code collision (less critical now without explicit nested transaction rollback on collision)
    # but still good to prevent duplicate active codes if possible, though the deactivate step helps.
    final_check_stmt = select(InviteModel.id).where(InviteModel.code == potential_code, InviteModel.is_active == True).limit(1)
    final_check_result = await db.execute(final_check_stmt)
    if final_check_result.scalar_one_or_none() is not None:
        # This is now more of a rare edge case if deactivate worked and code generation is diverse.
        # Depending on strictness, could raise an error or just log and proceed, 
        # relying on the previous deactivation to ensure only one is active.
        # For now, let's raise to be safe, as it implies a very quick collision.
        raise InviteOperationError("Invite code collision detected just before creation attempt.")

    db_invite = InviteModel(
        code=potential_code,
        group_id=group_id,
        created_by_id=creator_id,
        expires_at=expires_at,
        is_active=True
    )
    db.add(db_invite)
    await db.flush() # Flush to get ID for re-fetch and ensure it's in session before potential re-fetch.

    # Re-fetch with relationships
    stmt = (
        select(InviteModel)
        .where(InviteModel.id == db_invite.id)
        .options(
            selectinload(InviteModel.group),
            selectinload(InviteModel.creator)
        )
    )
    result = await db.execute(stmt)
    loaded_invite = result.scalar_one_or_none()

    if loaded_invite is None:
        # This would be an issue, implies flush didn't work or ID was wrong.
        # The main transaction will rollback if this exception is raised.
        raise InviteOperationError("Failed to load invite after creation and flush.")

    return loaded_invite
    # No explicit commit here, FastAPI handles it for the request.

async def get_active_invite_for_group(db: AsyncSession, group_id: int) -> Optional[InviteModel]:
    """Gets the currently active and non-expired invite for a specific group."""
    now = datetime.now(timezone.utc)
    try:
        stmt = (
            select(InviteModel).where(
                InviteModel.group_id == group_id,
                InviteModel.is_active == True,
                InviteModel.expires_at > now # Still respect expiry, even if very long
            )
            .order_by(InviteModel.created_at.desc()) # Get the most recent one if multiple (should not happen)
            .limit(1)
            .options(
                selectinload(InviteModel.group), # Eager load group
                selectinload(InviteModel.creator) # Eager load creator
            )
        )
        result = await db.execute(stmt)
        return result.scalars().first()
    except OperationalError as e:
        raise DatabaseConnectionError(f"DB connection error fetching active invite for group {group_id}: {str(e)}")
    except SQLAlchemyError as e:
        raise DatabaseQueryError(f"DB query error fetching active invite for group {group_id}: {str(e)}")

async def get_active_invite_by_code(db: AsyncSession, code: str) -> Optional[InviteModel]:
    """Gets an active and non-expired invite by its code."""
    now = datetime.now(timezone.utc)
    try:
        stmt = (
            select(InviteModel).where(
                InviteModel.code == code,
                InviteModel.is_active == True,
                InviteModel.expires_at > now
            )
            .options(
                selectinload(InviteModel.group),
                selectinload(InviteModel.creator)
            )
        )
        result = await db.execute(stmt)
        return result.scalars().first()
    except OperationalError as e:
        raise DatabaseConnectionError(f"DB connection error fetching invite: {str(e)}")
    except SQLAlchemyError as e:
        raise DatabaseQueryError(f"DB query error fetching invite: {str(e)}")

async def deactivate_invite(db: AsyncSession, invite: InviteModel) -> InviteModel:
    """Marks an invite as inactive (used) and reloads with relationships."""
    try:
        async with db.begin_nested() if db.in_transaction() else db.begin() as transaction:
            invite.is_active = False
            db.add(invite) # Add to session to track change
            await db.flush() # Persist is_active change

            # Re-fetch with relationships
            stmt = (
                select(InviteModel)
                .where(InviteModel.id == invite.id)
                .options(
                    selectinload(InviteModel.group),
                    selectinload(InviteModel.creator)
                )
            )
            result = await db.execute(stmt)
            updated_invite = result.scalar_one_or_none()

            if updated_invite is None: # Should not happen as invite is passed in
                await transaction.rollback()
                raise InviteOperationError("Failed to load invite after deactivation.")

            await transaction.commit()
            return updated_invite
    except OperationalError as e:
        raise DatabaseConnectionError(f"DB connection error deactivating invite: {str(e)}")
    except SQLAlchemyError as e:
        raise DatabaseTransactionError(f"DB transaction error deactivating invite: {str(e)}")

# Ensure InviteOperationError is defined in app.core.exceptions
# Example: class InviteOperationError(AppException): pass

# Optional: Function to periodically delete old, inactive invites
# async def cleanup_old_invites(db: AsyncSession, older_than_days: int = 30): ...