mitlist/be/app/crud/chore.py

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy.orm import selectinload, subqueryload
from sqlalchemy import union_all, and_, or_, delete
from typing import List, Optional
import logging
from datetime import date, datetime

from app.models import Chore, Group, User, ChoreAssignment, ChoreFrequencyEnum, ChoreTypeEnum, UserGroup, ChoreHistoryEventTypeEnum, UserRoleEnum, ChoreHistory, ChoreAssignmentHistory
from app.schemas.chore import ChoreCreate, ChoreUpdate, ChoreAssignmentCreate, ChoreAssignmentUpdate
from app.core.chore_utils import calculate_next_due_date
from app.crud.group import get_group_by_id, is_user_member, get_user_role_in_group
from app.crud.history import create_chore_history_entry, create_assignment_history_entry
from app.core.exceptions import ChoreNotFoundError, GroupNotFoundError, PermissionDeniedError, DatabaseIntegrityError

logger = logging.getLogger(__name__)

async def get_all_user_chores(db: AsyncSession, user_id: int) -> List[Chore]:
    """Gets all chores (personal and group) for a user in optimized queries."""
    
    personal_chores_query = (
        select(Chore)   
        .where(
            Chore.created_by_id == user_id,
            Chore.type == ChoreTypeEnum.personal
        )
    )
    
    user_groups_result = await db.execute(
        select(UserGroup.group_id).where(UserGroup.user_id == user_id)
    )
    user_group_ids = user_groups_result.scalars().all()
    
    all_chores = []
    
    personal_result = await db.execute(
        personal_chores_query
        .options(
            selectinload(Chore.creator),
            selectinload(Chore.assignments).selectinload(ChoreAssignment.assigned_user),
            selectinload(Chore.assignments).selectinload(ChoreAssignment.history),
            selectinload(Chore.history),
            selectinload(Chore.child_chores)
        )
        .order_by(Chore.next_due_date, Chore.name)
    )
    all_chores.extend(personal_result.scalars().all())
    
    if user_group_ids:
        group_chores_result = await db.execute(
            select(Chore)
            .where(
                Chore.group_id.in_(user_group_ids),
                Chore.type == ChoreTypeEnum.group
            )
            .options(
                selectinload(Chore.creator),
                selectinload(Chore.group),
                selectinload(Chore.assignments).selectinload(ChoreAssignment.assigned_user),
                selectinload(Chore.assignments).selectinload(ChoreAssignment.history),
                selectinload(Chore.history),
                selectinload(Chore.child_chores)
            )
            .order_by(Chore.next_due_date, Chore.name)
        )
        all_chores.extend(group_chores_result.scalars().all())
    
    return all_chores

async def create_chore(
    db: AsyncSession,
    chore_in: ChoreCreate,
    user_id: int
) -> Chore:
    """Creates a new chore, and if specified, an assignment for it."""
    async with db.begin_nested() if db.in_transaction() else db.begin():
        # Validate chore type and group
        if chore_in.type == ChoreTypeEnum.group:
            if not chore_in.group_id:
                raise ValueError("group_id is required for group chores")
            group = await get_group_by_id(db, chore_in.group_id)
            if not group:
                raise GroupNotFoundError(chore_in.group_id)
            if not await is_user_member(db, chore_in.group_id, user_id):
                raise PermissionDeniedError(detail=f"User {user_id} not a member of group {chore_in.group_id}")
        else:  # personal chore
            if chore_in.group_id:
                raise ValueError("group_id must be None for personal chores")

        # Validate assigned user if provided
        if chore_in.assigned_to_user_id:
            if chore_in.type == ChoreTypeEnum.group:
                # For group chores, assigned user must be a member of the group
                if not await is_user_member(db, chore_in.group_id, chore_in.assigned_to_user_id):
                    raise PermissionDeniedError(detail=f"Assigned user {chore_in.assigned_to_user_id} is not a member of group {chore_in.group_id}")
            else: # Personal chore
                # For personal chores, you can only assign it to yourself
                if chore_in.assigned_to_user_id != user_id:
                    raise PermissionDeniedError(detail="Personal chores can only be assigned to the creator.")
        
        assigned_user_id = chore_in.assigned_to_user_id
        chore_data = chore_in.model_dump(exclude_unset=True, exclude={'assigned_to_user_id'})

        if 'parent_chore_id' in chore_data and chore_data['parent_chore_id']:
            parent_chore = await get_chore_by_id(db, chore_data['parent_chore_id'])
            if not parent_chore:
                raise ChoreNotFoundError(chore_data['parent_chore_id'])

        db_chore = Chore(
            **chore_data,
            created_by_id=user_id,
        )

        if chore_in.frequency == ChoreFrequencyEnum.custom and chore_in.custom_interval_days is None:
            raise ValueError("custom_interval_days must be set for custom frequency chores.")

        db.add(db_chore)
        await db.flush()
        
        # Create an assignment if a user was specified
        if assigned_user_id:
            assignment = ChoreAssignment(
                chore_id=db_chore.id,
                assigned_to_user_id=assigned_user_id,
                due_date=db_chore.next_due_date,
                is_complete=False
            )
            db.add(assignment)
            await db.flush()  # Flush to get the assignment ID
            await create_assignment_history_entry(
                db, 
                assignment_id=assignment.id, 
                event_type=ChoreHistoryEventTypeEnum.ASSIGNED, 
                changed_by_user_id=user_id,
                event_data={'assigned_to': assigned_user_id}
            )

        await create_chore_history_entry(
            db,
            chore_id=db_chore.id,
            group_id=db_chore.group_id,
            changed_by_user_id=user_id,
            event_type=ChoreHistoryEventTypeEnum.CREATED,
            event_data={"chore_name": db_chore.name}
        )

        try:
            result = await db.execute(
                select(Chore)
                .where(Chore.id == db_chore.id)
                .options(
                    selectinload(Chore.creator), 
                    selectinload(Chore.group), 
                    selectinload(Chore.assignments).selectinload(ChoreAssignment.assigned_user),
                    selectinload(Chore.assignments).selectinload(ChoreAssignment.history),
                    selectinload(Chore.history),
                    selectinload(Chore.child_chores)
                )
            )
            return result.scalar_one()
        except Exception as e:
            logger.error(f"Error creating chore: {e}", exc_info=True)
            raise DatabaseIntegrityError(f"Could not create chore. Error: {str(e)}")

async def get_chore_by_id(db: AsyncSession, chore_id: int) -> Optional[Chore]:
    """Gets a chore by ID."""
    result = await db.execute(
        select(Chore)
        .where(Chore.id == chore_id)
        .options(*get_chore_loader_options())
    )
    return result.scalar_one_or_none()

async def get_chore_by_id_and_group(
    db: AsyncSession, 
    chore_id: int, 
    group_id: int, 
    user_id: int
) -> Optional[Chore]:
    """Gets a specific group chore by ID, ensuring it belongs to the group and user is a member."""
    if not await is_user_member(db, group_id, user_id):
        raise PermissionDeniedError(detail=f"User {user_id} not a member of group {group_id}")
    
    chore = await get_chore_by_id(db, chore_id)
    if chore and chore.group_id == group_id and chore.type == ChoreTypeEnum.group:
        return chore
    return None

async def get_personal_chores(
    db: AsyncSession,
    user_id: int
) -> List[Chore]:
    """Gets all personal chores for a user with optimized eager loading."""
    result = await db.execute(
        select(Chore)
        .where(
            Chore.created_by_id == user_id,
            Chore.type == ChoreTypeEnum.personal
        )
        .options(*get_chore_loader_options())
        .order_by(Chore.next_due_date, Chore.name)
    )
    return result.scalars().all()

async def get_chores_by_group_id(
    db: AsyncSession, 
    group_id: int, 
    user_id: int
) -> List[Chore]:
    """Gets all chores for a specific group with optimized eager loading, if the user is a member."""
    if not await is_user_member(db, group_id, user_id):
        raise PermissionDeniedError(detail=f"User {user_id} not a member of group {group_id}")

    result = await db.execute(
        select(Chore)
        .where(
            Chore.group_id == group_id,
            Chore.type == ChoreTypeEnum.group
        )
        .options(*get_chore_loader_options())
        .order_by(Chore.next_due_date, Chore.name)
    )
    return result.scalars().all()

async def update_chore(
    db: AsyncSession,
    chore_id: int,
    chore_in: ChoreUpdate,
    user_id: int,
    group_id: Optional[int] = None
) -> Optional[Chore]:
    """Updates a chore's details using proper transaction management."""
    async with db.begin_nested() if db.in_transaction() else db.begin():
        db_chore = await get_chore_by_id(db, chore_id)
        if not db_chore:
            raise ChoreNotFoundError(chore_id, group_id)

        original_data = {field: getattr(db_chore, field) for field in chore_in.model_dump(exclude_unset=True)}

        # Check permissions for current chore
        if db_chore.type == ChoreTypeEnum.group:
            if not await is_user_member(db, db_chore.group_id, user_id):
                raise PermissionDeniedError(detail=f"User {user_id} not a member of chore's current group {db_chore.group_id}")
        else: 
            if db_chore.created_by_id != user_id:
                raise PermissionDeniedError(detail="Only the creator can update personal chores")

        update_data = chore_in.model_dump(exclude_unset=True)
        
        # Handle group changes
        if 'group_id' in update_data:
            new_group_id = update_data['group_id']
            if new_group_id != db_chore.group_id:
                # Validate user has permission for the new group
                if new_group_id is not None:
                    if not await is_user_member(db, new_group_id, user_id):
                        raise PermissionDeniedError(detail=f"User {user_id} not a member of target group {new_group_id}")
        
        # Handle type changes
        if 'type' in update_data:
            new_type = update_data['type']
            # When changing to personal, always clear group_id regardless of what's in update_data
            if new_type == ChoreTypeEnum.personal:
                update_data['group_id'] = None
            else:
                # For group chores, use the provided group_id or keep the current one
                new_group_id = update_data.get('group_id', db_chore.group_id)
                if new_type == ChoreTypeEnum.group and new_group_id is None:
                    raise ValueError("group_id is required for group chores")
        
        if 'parent_chore_id' in update_data:
            if update_data['parent_chore_id']:
                parent_chore = await get_chore_by_id(db, update_data['parent_chore_id'])
                if not parent_chore:
                    raise ChoreNotFoundError(update_data['parent_chore_id'])
            # Setting parent_chore_id to None is allowed
            setattr(db_chore, 'parent_chore_id', update_data['parent_chore_id'])

        # Recalculate next_due_date if needed
        recalculate = False
        if 'frequency' in update_data and update_data['frequency'] != db_chore.frequency:
            recalculate = True
        if 'custom_interval_days' in update_data and update_data['custom_interval_days'] != db_chore.custom_interval_days:
            recalculate = True

        current_next_due_date_for_calc = db_chore.next_due_date
        if 'next_due_date' in update_data:
            current_next_due_date_for_calc = update_data['next_due_date']
            if not ('frequency' in update_data or 'custom_interval_days' in update_data):
                recalculate = False

        for field, value in update_data.items():
            setattr(db_chore, field, value)
        
        if recalculate:
            db_chore.next_due_date = calculate_next_due_date(
                current_due_date=current_next_due_date_for_calc,
                frequency=db_chore.frequency,
                custom_interval_days=db_chore.custom_interval_days,
                last_completed_date=db_chore.last_completed_at
            )
        
        if db_chore.frequency == ChoreFrequencyEnum.custom and db_chore.custom_interval_days is None:
            raise ValueError("custom_interval_days must be set for custom frequency chores.")

        changes = {}
        for field, old_value in original_data.items():
            new_value = getattr(db_chore, field)
            if old_value != new_value:
                changes[field] = {"old": str(old_value), "new": str(new_value)}
        
        if changes:
            await create_chore_history_entry(
                db,
                chore_id=chore_id,
                group_id=db_chore.group_id,
                changed_by_user_id=user_id,
                event_type=ChoreHistoryEventTypeEnum.UPDATED,
                event_data=changes
            )

        try:
            await db.flush()
            result = await db.execute(
                select(Chore)
                .where(Chore.id == db_chore.id)
                .options(*get_chore_loader_options())
            )
            return result.scalar_one()
        except Exception as e:
            logger.error(f"Error updating chore {chore_id}: {e}", exc_info=True)
            raise DatabaseIntegrityError(f"Could not update chore {chore_id}. Error: {str(e)}")

async def delete_chore(
    db: AsyncSession,
    chore_id: int,
    user_id: int,
    group_id: Optional[int] = None
) -> bool:
    """Deletes a chore and its assignments using proper transaction management, ensuring user has permission."""
    async with db.begin_nested() if db.in_transaction() else db.begin():
        db_chore = await get_chore_by_id(db, chore_id)
        if not db_chore:
            raise ChoreNotFoundError(chore_id, group_id)

        await create_chore_history_entry(
            db,
            chore_id=chore_id,
            group_id=db_chore.group_id,
            changed_by_user_id=user_id,
            event_type=ChoreHistoryEventTypeEnum.DELETED,
            event_data={"chore_name": db_chore.name}
        )

        if db_chore.type == ChoreTypeEnum.group:
            if not group_id:
                raise ValueError("group_id is required for group chores")
            if not await is_user_member(db, group_id, user_id):
                raise PermissionDeniedError(detail=f"User {user_id} not a member of group {group_id}")
            if db_chore.group_id != group_id:
                raise ChoreNotFoundError(chore_id, group_id)
        else:  # personal chore
            if group_id:
                raise ValueError("group_id must be None for personal chores")
            if db_chore.created_by_id != user_id:
                raise PermissionDeniedError(detail="Only the creator can delete personal chores")

        try:
            await db.delete(db_chore)
            await db.flush()
            return True
        except Exception as e:
            logger.error(f"Error deleting chore {chore_id}: {e}", exc_info=True)
            raise DatabaseIntegrityError(f"Could not delete chore {chore_id}. Error: {str(e)}")

# === CHORE ASSIGNMENT CRUD FUNCTIONS ===

async def create_chore_assignment(
    db: AsyncSession,
    assignment_in: ChoreAssignmentCreate,
    user_id: int
) -> ChoreAssignment:
    """Creates a new chore assignment. User must be able to manage the chore."""
    async with db.begin_nested() if db.in_transaction() else db.begin():
        chore = await get_chore_by_id(db, assignment_in.chore_id)
        if not chore:
            raise ChoreNotFoundError(chore_id=assignment_in.chore_id)
        
        if chore.type == ChoreTypeEnum.personal:
            if chore.created_by_id != user_id:
                raise PermissionDeniedError(detail="Only the creator can assign personal chores")
        else:  # group chore
            if not await is_user_member(db, chore.group_id, user_id):
                raise PermissionDeniedError(detail=f"User {user_id} not a member of group {chore.group_id}")
            if not await is_user_member(db, chore.group_id, assignment_in.assigned_to_user_id):
                raise PermissionDeniedError(detail=f"Cannot assign chore to user {assignment_in.assigned_to_user_id} who is not a group member")

        db_assignment = ChoreAssignment(**assignment_in.model_dump(exclude_unset=True))
        db.add(db_assignment)
        await db.flush()
        
        await create_assignment_history_entry(
            db,
            assignment_id=db_assignment.id,
            changed_by_user_id=user_id,
            event_type=ChoreHistoryEventTypeEnum.ASSIGNED,
            event_data={"assigned_to_user_id": db_assignment.assigned_to_user_id, "due_date": db_assignment.due_date.isoformat()}
        )
        
        try:
            result = await db.execute(
                select(ChoreAssignment)
                .where(ChoreAssignment.id == db_assignment.id)
                .options(*get_assignment_loader_options())
            )
            return result.scalar_one()
        except Exception as e:
            logger.error(f"Error creating chore assignment: {e}", exc_info=True)
            raise DatabaseIntegrityError(f"Could not create chore assignment. Error: {str(e)}")

async def get_chore_assignment_by_id(db: AsyncSession, assignment_id: int) -> Optional[ChoreAssignment]:
    """Gets a chore assignment by ID."""
    result = await db.execute(
        select(ChoreAssignment)
        .where(ChoreAssignment.id == assignment_id)
        .options(*get_assignment_loader_options())
    )
    return result.scalar_one_or_none()

async def get_user_assignments(
    db: AsyncSession,
    user_id: int,
    include_completed: bool = False
) -> List[ChoreAssignment]:
    """Gets all chore assignments for a user."""
    query = select(ChoreAssignment).where(ChoreAssignment.assigned_to_user_id == user_id)
    
    if not include_completed:
        query = query.where(ChoreAssignment.is_complete == False)
    
    query = query.options(*get_assignment_loader_options()).order_by(ChoreAssignment.due_date, ChoreAssignment.id)
    
    result = await db.execute(query)
    return result.scalars().all()

async def get_chore_assignments(
    db: AsyncSession,
    chore_id: int,
    user_id: int
) -> List[ChoreAssignment]:
    """Gets all assignments for a specific chore. User must have permission to view the chore."""
    chore = await get_chore_by_id(db, chore_id)
    if not chore:
        raise ChoreNotFoundError(chore_id=chore_id)

    if chore.type == ChoreTypeEnum.personal:
        if chore.created_by_id != user_id:
            raise PermissionDeniedError(detail="Can only view assignments for own personal chores")
    else:
        if not await is_user_member(db, chore.group_id, user_id):
            raise PermissionDeniedError(detail=f"User {user_id} not a member of group {chore.group_id}")
    
    result = await db.execute(
        select(ChoreAssignment)
        .where(ChoreAssignment.chore_id == chore_id)
        .options(*get_assignment_loader_options())
        .order_by(ChoreAssignment.due_date, ChoreAssignment.id)
    )
    return result.scalars().all()

async def update_chore_assignment(
    db: AsyncSession,
    assignment_id: int,
    assignment_in: ChoreAssignmentUpdate,
    user_id: int
) -> Optional[ChoreAssignment]:
    """Updates a chore assignment, e.g., to mark it as complete."""
    async with db.begin_nested() if db.in_transaction() else db.begin():
        db_assignment = await get_chore_assignment_by_id(db, assignment_id)
        if not db_assignment:
            return None

        # Permission Check: only assigned user or group owner can update
        is_allowed = db_assignment.assigned_to_user_id == user_id
        if not is_allowed and db_assignment.chore.group_id:
            user_role = await get_user_role_in_group(db, db_assignment.chore.group_id, user_id)
            is_allowed = user_role == UserRoleEnum.owner
        
        if not is_allowed:
            raise PermissionDeniedError("You cannot update this chore assignment.")
        
        original_status = db_assignment.is_complete
        update_data = assignment_in.model_dump(exclude_unset=True)

        for field, value in update_data.items():
            setattr(db_assignment, field, value)

        if 'is_complete' in update_data:
            new_status = update_data['is_complete']
            history_event = None
            if new_status and not original_status:
                db_assignment.completed_at = datetime.utcnow()
                history_event = ChoreHistoryEventTypeEnum.COMPLETED
                
                # Advance the next_due_date of the parent chore
                if db_assignment.chore and db_assignment.chore.frequency != ChoreFrequencyEnum.one_time:
                    db_assignment.chore.last_completed_at = db_assignment.completed_at
                    db_assignment.chore.next_due_date = calculate_next_due_date(
                        current_due_date=db_assignment.chore.next_due_date,
                        frequency=db_assignment.chore.frequency,
                        custom_interval_days=db_assignment.chore.custom_interval_days,
                        last_completed_date=db_assignment.chore.last_completed_at.date() if db_assignment.chore.last_completed_at else None
                    )
            elif not new_status and original_status:
                db_assignment.completed_at = None
                history_event = ChoreHistoryEventTypeEnum.REOPENED
                # Policy: Do not automatically roll back parent chore's due date.

            if history_event:
                await create_assignment_history_entry(
                    db=db,
                    assignment_id=assignment_id,
                    changed_by_user_id=user_id,
                    event_type=history_event,
                    event_data={"new_status": new_status}
                )

        await db.flush()
        
        try:
            result = await db.execute(
                select(ChoreAssignment)
                .where(ChoreAssignment.id == assignment_id)
                .options(*get_assignment_loader_options())
            )
            return result.scalar_one()
        except Exception as e:
            logger.error(f"Error updating assignment: {e}", exc_info=True)
            await db.rollback()
            raise DatabaseIntegrityError(f"Could not update assignment. Error: {str(e)}")

async def delete_chore_assignment(
    db: AsyncSession,
    assignment_id: int,
    user_id: int
) -> bool:
    """Deletes a chore assignment. User must have permission to manage the chore."""
    async with db.begin_nested() if db.in_transaction() else db.begin():
        db_assignment = await get_chore_assignment_by_id(db, assignment_id)
        if not db_assignment:
            raise ChoreNotFoundError(assignment_id=assignment_id)
        
        await create_assignment_history_entry(
            db,
            assignment_id=assignment_id,
            changed_by_user_id=user_id,
            event_type=ChoreHistoryEventTypeEnum.UNASSIGNED,
            event_data={"unassigned_user_id": db_assignment.assigned_to_user_id}
        )

        chore = await get_chore_by_id(db, db_assignment.chore_id)
        if not chore:
            raise ChoreNotFoundError(chore_id=db_assignment.chore_id)
        
        if chore.type == ChoreTypeEnum.personal:
            if chore.created_by_id != user_id:
                raise PermissionDeniedError(detail="Only the creator can delete personal chore assignments")
        else:
            if not await is_user_member(db, chore.group_id, user_id):
                raise PermissionDeniedError(detail=f"User {user_id} not a member of group {chore.group_id}")
        
        try:
            await db.delete(db_assignment)
            await db.flush()
            return True
        except Exception as e:
            logger.error(f"Error deleting chore assignment {assignment_id}: {e}", exc_info=True)
            raise DatabaseIntegrityError(f"Could not delete chore assignment {assignment_id}. Error: {str(e)}")

def get_chore_loader_options():
    """Returns a list of SQLAlchemy loader options for chore relationships."""
    return [
        selectinload(Chore.creator),
        selectinload(Chore.group),
        selectinload(Chore.assignments).selectinload(ChoreAssignment.assigned_user),
        selectinload(Chore.history).selectinload(ChoreHistory.changed_by_user),
        selectinload(Chore.child_chores).options(
            selectinload(Chore.assignments).selectinload(ChoreAssignment.assigned_user),
            selectinload(Chore.history).selectinload(ChoreHistory.changed_by_user),
            selectinload(Chore.creator),
            selectinload(Chore.child_chores)  # Load grandchildren, adjust depth if needed
        )
    ]

def get_assignment_loader_options():
    return [
        selectinload(ChoreAssignment.assigned_user),
        selectinload(ChoreAssignment.history).selectinload(ChoreAssignmentHistory.changed_by_user),
        selectinload(ChoreAssignment.chore).options(*get_chore_loader_options())
    ]