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

from app.models import Chore, Group, User, ChoreAssignment, ChoreFrequencyEnum, ChoreTypeEnum, UserGroup
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
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."""
    
    # Get personal chores query
    personal_chores_query = (
        select(Chore)
        .where(
            Chore.created_by_id == user_id,
            Chore.type == ChoreTypeEnum.personal
        )
    )
    
    # Get user's group IDs first
    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 = []
    
    # Execute personal chores query
    personal_result = await db.execute(
        personal_chores_query
        .options(
            selectinload(Chore.creator),
            selectinload(Chore.assignments).selectinload(ChoreAssignment.assigned_user)
        )
        .order_by(Chore.next_due_date, Chore.name)
    )
    all_chores.extend(personal_result.scalars().all())
    
    # If user has groups, get all group chores in one query
    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)
            )
            .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,
    group_id: Optional[int] = None
) -> Chore:
    """Creates a new chore, either personal or within a specific group."""
    # Use the transaction pattern from the FastAPI strategy
    async with db.begin_nested() if db.in_transaction() else db.begin():
        if chore_in.type == ChoreTypeEnum.group:
            if not group_id:
                raise ValueError("group_id is required for group chores")
            # Validate group existence and user membership
            group = await get_group_by_id(db, group_id)
            if not group:
                raise GroupNotFoundError(group_id)
            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}")
        else:  # personal chore
            if group_id:
                raise ValueError("group_id must be None for personal chores")

        db_chore = Chore(
            **chore_in.model_dump(exclude_unset=True, exclude={'group_id'}),
            group_id=group_id,
            created_by_id=user_id,
        )

        # Specific check for custom frequency
        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()  # Get the ID for the chore
        
        try:
            # Load relationships for the response with eager loading
            result = await db.execute(
                select(Chore)
                .where(Chore.id == db_chore.id)
                .options(
                    selectinload(Chore.creator), 
                    selectinload(Chore.group), 
                    selectinload(Chore.assignments)
                )
            )
            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(selectinload(Chore.creator), selectinload(Chore.group), selectinload(Chore.assignments))
    )
    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(
            selectinload(Chore.creator),
            selectinload(Chore.assignments).selectinload(ChoreAssignment.assigned_user)
        )
        .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(
            selectinload(Chore.creator),
            selectinload(Chore.assignments).selectinload(ChoreAssignment.assigned_user)
        )
        .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)

        # Check permissions
        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 update personal chores")

        update_data = chore_in.model_dump(exclude_unset=True)

        # Handle type change
        if 'type' in update_data:
            new_type = update_data['type']
            if new_type == ChoreTypeEnum.group and not group_id:
                raise ValueError("group_id is required for group chores")
            if new_type == ChoreTypeEnum.personal and group_id:
                raise ValueError("group_id must be None for personal chores")

        # 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.")

        try:
            await db.flush()  # Flush changes within the transaction
            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)
                )
            )
            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)

        # Check permissions
        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()  # Ensure deletion is processed within the transaction
            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():
        # Get the chore and validate permissions
        chore = await get_chore_by_id(db, assignment_in.chore_id)
        if not chore:
            raise ChoreNotFoundError(chore_id=assignment_in.chore_id)
        
        # Check permissions to assign this chore
        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}")
            # For group chores, check if assignee is also a group member
            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()  # Get the ID for the assignment
        
        try:
            # Load relationships for the response
            result = await db.execute(
                select(ChoreAssignment)
                .where(ChoreAssignment.id == db_assignment.id)
                .options(
                    selectinload(ChoreAssignment.chore).selectinload(Chore.creator),
                    selectinload(ChoreAssignment.assigned_user)
                )
            )
            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(
            selectinload(ChoreAssignment.chore).selectinload(Chore.creator),
            selectinload(ChoreAssignment.assigned_user)
        )
    )
    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(
        selectinload(ChoreAssignment.chore).selectinload(Chore.creator),
        selectinload(ChoreAssignment.assigned_user)
    ).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)
    
    # Check permissions
    if chore.type == ChoreTypeEnum.personal:
        if chore.created_by_id != user_id:
            raise PermissionDeniedError(detail="Can only view assignments for own 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}")
    
    result = await db.execute(
        select(ChoreAssignment)
        .where(ChoreAssignment.chore_id == chore_id)
        .options(
            selectinload(ChoreAssignment.chore).selectinload(Chore.creator),
            selectinload(ChoreAssignment.assigned_user)
        )
        .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. Only the assignee can mark it 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:
            raise ChoreNotFoundError(assignment_id=assignment_id)
        
        # Load the chore for permission checking
        chore = await get_chore_by_id(db, db_assignment.chore_id)
        if not chore:
            raise ChoreNotFoundError(chore_id=db_assignment.chore_id)
        
        # Check permissions - only assignee can complete, but chore managers can reschedule
        can_manage = False
        if chore.type == ChoreTypeEnum.personal:
            can_manage = chore.created_by_id == user_id
        else:  # group chore
            can_manage = await is_user_member(db, chore.group_id, user_id)
        
        can_complete = db_assignment.assigned_to_user_id == user_id
        
        update_data = assignment_in.model_dump(exclude_unset=True)
        
        # Check specific permissions for different updates
        if 'is_complete' in update_data and not can_complete:
            raise PermissionDeniedError(detail="Only the assignee can mark assignments as complete")
        
        if 'due_date' in update_data and not can_manage:
            raise PermissionDeniedError(detail="Only chore managers can reschedule assignments")
        
        # Handle completion logic
        if 'is_complete' in update_data and update_data['is_complete']:
            if not db_assignment.is_complete:  # Only if not already complete
                update_data['completed_at'] = datetime.utcnow()
                
                # Update parent chore's last_completed_at and recalculate next_due_date
                chore.last_completed_at = update_data['completed_at']
                chore.next_due_date = calculate_next_due_date(
                    current_due_date=chore.next_due_date,
                    frequency=chore.frequency,
                    custom_interval_days=chore.custom_interval_days,
                    last_completed_date=chore.last_completed_at
                )
        elif 'is_complete' in update_data and not update_data['is_complete']:
            # If marking as incomplete, clear completed_at
            update_data['completed_at'] = None
        
        # Apply updates
        for field, value in update_data.items():
            setattr(db_assignment, field, value)
        
        try:
            await db.flush()  # Flush changes within the transaction
            
            # Load relationships for the response
            result = await db.execute(
                select(ChoreAssignment)
                .where(ChoreAssignment.id == db_assignment.id)
                .options(
                    selectinload(ChoreAssignment.chore).selectinload(Chore.creator),
                    selectinload(ChoreAssignment.assigned_user)
                )
            )
            return result.scalar_one()
        except Exception as e:
            logger.error(f"Error updating chore assignment {assignment_id}: {e}", exc_info=True)
            raise DatabaseIntegrityError(f"Could not update chore assignment {assignment_id}. 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)
        
        # Load the chore for permission checking
        chore = await get_chore_by_id(db, db_assignment.chore_id)
        if not chore:
            raise ChoreNotFoundError(chore_id=db_assignment.chore_id)
        
        # Check permissions
        if chore.type == ChoreTypeEnum.personal:
            if chore.created_by_id != user_id:
                raise PermissionDeniedError(detail="Only the creator can delete personal chore assignments")
        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}")
        
        try:
            await db.delete(db_assignment)
            await db.flush()  # Ensure deletion is processed within the transaction
            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)}")