"""Workload calculation service. Provides functionality to calculate and retrieve user workload data including weekly load percentages, task allocations, and load level classification. """ from datetime import date, timedelta from decimal import Decimal from typing import List, Optional, Tuple from sqlalchemy import func, and_ from sqlalchemy.orm import Session, joinedload from app.models.user import User from app.models.task import Task from app.models.task_status import TaskStatus from app.models.project import Project from app.schemas.workload import ( LoadLevel, UserWorkloadSummary, UserWorkloadDetail, TaskWorkloadInfo, ) def get_week_bounds(d: date) -> Tuple[date, date]: """ Get ISO week boundaries (Monday to Sunday). Args: d: Any date within the week Returns: Tuple of (week_start, week_end) where week_start is Monday """ week_start = d - timedelta(days=d.weekday()) week_end = week_start + timedelta(days=6) return week_start, week_end def get_current_week_start() -> date: """Get the Monday of the current week.""" return get_week_bounds(date.today())[0] def get_current_week_bounds() -> Tuple[date, date]: """ Get current week bounds for default views. On Sundays, extend the window to include the upcoming week so that "tomorrow" tasks are still visible in default views. """ week_start, week_end = get_week_bounds(date.today()) if date.today().weekday() == 6: week_end = week_end + timedelta(days=7) return week_start, week_end def _extend_week_end_if_sunday(week_start: date, week_end: date) -> Tuple[date, date]: """Extend week window on Sunday to include upcoming week.""" if date.today().weekday() == 6 and week_start == get_current_week_start(): return week_start, week_end + timedelta(days=7) return week_start, week_end def determine_load_level(load_percentage: Optional[Decimal]) -> LoadLevel: """ Determine the load level based on percentage. Args: load_percentage: The calculated load percentage (None if capacity is 0) Returns: LoadLevel enum value """ if load_percentage is None: return LoadLevel.UNAVAILABLE if load_percentage < 80: return LoadLevel.NORMAL elif load_percentage < 100: return LoadLevel.WARNING else: return LoadLevel.OVERLOADED def calculate_load_percentage( allocated_hours: Decimal, capacity_hours: Decimal ) -> Optional[Decimal]: """ Calculate load percentage avoiding division by zero. Args: allocated_hours: Total allocated hours capacity_hours: User's weekly capacity Returns: Load percentage or None if capacity is 0 """ if capacity_hours == 0: return None return (allocated_hours / capacity_hours * 100).quantize(Decimal("0.01")) def get_user_tasks_in_week( db: Session, user_id: str, week_start: date, week_end: date, ) -> List[Task]: """ Get all tasks assigned to a user with due_date in the specified week. Excludes tasks with is_done=True status. Args: db: Database session user_id: User ID week_start: Start of week (Monday) week_end: End of week (Sunday) Returns: List of Task objects """ # Convert date to datetime for comparison from datetime import datetime week_start_dt = datetime.combine(week_start, datetime.min.time()) week_end_dt = datetime.combine(week_end, datetime.max.time()) return ( db.query(Task) .join(Task.status, isouter=True) .join(Task.project) .filter( Task.assignee_id == user_id, Task.due_date >= week_start_dt, Task.due_date <= week_end_dt, # Exclude completed tasks (TaskStatus.is_done == False) | (Task.status_id == None) ) .options(joinedload(Task.project), joinedload(Task.status)) .all() ) def calculate_user_workload( db: Session, user: User, week_start: date, ) -> UserWorkloadSummary: """ Calculate workload summary for a single user. Args: db: Database session user: User object week_start: Start of week (Monday) Returns: UserWorkloadSummary object """ week_start, week_end = get_week_bounds(week_start) # Get tasks for this user in this week tasks = get_user_tasks_in_week(db, user.id, week_start, week_end) # Calculate allocated hours from original_estimate allocated_hours = Decimal("0") for task in tasks: if task.original_estimate: allocated_hours += task.original_estimate capacity_hours = Decimal(str(user.capacity)) if user.capacity is not None else Decimal("40") load_percentage = calculate_load_percentage(allocated_hours, capacity_hours) load_level = determine_load_level(load_percentage) return UserWorkloadSummary( user_id=user.id, user_name=user.name, department_id=user.department_id, department_name=user.department.name if user.department else None, capacity_hours=capacity_hours, allocated_hours=allocated_hours, load_percentage=load_percentage, load_level=load_level, task_count=len(tasks), ) def get_workload_heatmap( db: Session, week_start: Optional[date] = None, department_id: Optional[str] = None, user_ids: Optional[List[str]] = None, hide_empty: bool = True, ) -> List[UserWorkloadSummary]: """ Get workload heatmap for multiple users. Args: db: Database session week_start: Start of week (defaults to current week) department_id: Filter by department user_ids: Filter by specific user IDs hide_empty: If True, exclude users with no tasks (default: True) Returns: List of UserWorkloadSummary objects """ from datetime import datetime from collections import defaultdict if week_start is None: week_start = get_current_week_start() # Normalize to week start (Monday) week_start = get_week_bounds(week_start)[0] week_start, week_end = get_week_bounds(week_start) week_start, week_end = _extend_week_end_if_sunday(week_start, week_end) # Build user query query = db.query(User).filter(User.is_active == True) if department_id: query = query.filter(User.department_id == department_id) if user_ids: query = query.filter(User.id.in_(user_ids)) users = query.options(joinedload(User.department)).all() if not users: return [] # Batch query: fetch all tasks for all users in one query user_id_list = [user.id for user in users] week_start_dt = datetime.combine(week_start, datetime.min.time()) week_end_dt = datetime.combine(week_end, datetime.max.time()) all_tasks = ( db.query(Task) .join(Task.status, isouter=True) .filter( Task.assignee_id.in_(user_id_list), Task.due_date >= week_start_dt, Task.due_date <= week_end_dt, # Exclude completed tasks (TaskStatus.is_done == False) | (Task.status_id == None) ) .all() ) # Group tasks by assignee_id in memory tasks_by_user: dict = defaultdict(list) for task in all_tasks: tasks_by_user[task.assignee_id].append(task) # Calculate workload for each user using pre-fetched tasks results = [] for user in users: user_tasks = tasks_by_user.get(user.id, []) # Calculate allocated hours from original_estimate allocated_hours = Decimal("0") for task in user_tasks: if task.original_estimate: allocated_hours += task.original_estimate capacity_hours = Decimal(str(user.capacity)) if user.capacity is not None else Decimal("40") load_percentage = calculate_load_percentage(allocated_hours, capacity_hours) load_level = determine_load_level(load_percentage) summary = UserWorkloadSummary( user_id=user.id, user_name=user.name, department_id=user.department_id, department_name=user.department.name if user.department else None, capacity_hours=capacity_hours, allocated_hours=allocated_hours, load_percentage=load_percentage, load_level=load_level, task_count=len(user_tasks), ) results.append(summary) # Filter out users with no tasks if hide_empty is True if hide_empty: results = [r for r in results if r.task_count > 0] return results def get_user_workload_detail( db: Session, user_id: str, week_start: Optional[date] = None, ) -> Optional[UserWorkloadDetail]: """ Get detailed workload for a specific user including task list. Args: db: Database session user_id: User ID week_start: Start of week (defaults to current week) Returns: UserWorkloadDetail object or None if user not found """ user = ( db.query(User) .filter(User.id == user_id) .options(joinedload(User.department)) .first() ) if not user: return None if week_start is None: week_start = get_current_week_start() week_start = get_week_bounds(week_start)[0] week_start, week_end = get_week_bounds(week_start) week_start, week_end = _extend_week_end_if_sunday(week_start, week_end) # Get tasks tasks = get_user_tasks_in_week(db, user_id, week_start, week_end) # Calculate totals allocated_hours = Decimal("0") task_infos = [] for task in tasks: if task.original_estimate: allocated_hours += task.original_estimate task_infos.append(TaskWorkloadInfo( task_id=task.id, title=task.title, project_id=task.project_id, project_name=task.project.title if task.project else "Unknown", due_date=task.due_date, original_estimate=task.original_estimate, status=task.status.name if task.status else None, )) capacity_hours = Decimal(str(user.capacity)) if user.capacity is not None else Decimal("40") load_percentage = calculate_load_percentage(allocated_hours, capacity_hours) load_level = determine_load_level(load_percentage) return UserWorkloadDetail( user_id=user.id, user_name=user.name, week_start=week_start, week_end=week_end, capacity_hours=capacity_hours, allocated_hours=allocated_hours, load_percentage=load_percentage, load_level=load_level, tasks=task_infos, )