dow-prod-tracker/src/lib/services/workload-service.ts

import { prisma } from "@/lib/prisma";
import {
  startOfWeek,
  endOfWeek,
  addWeeks,
  format,
  isWithinInterval,
} from "date-fns";
import type { VisibilityContext } from "@/lib/rbac/visibility";

/** Statuses that count as "active" work (not terminal). */
const ACTIVE_STATUSES = [
  "NOT_STARTED",
  "IN_PROGRESS",
  "IN_REVIEW",
  "CHANGES_REQUESTED",
] as const;

/** Statuses that mean the stage is complete. */
const TERMINAL_STATUSES = [
  "APPROVED",
  "DELIVERED",
  "SKIPPED",
] as const;

export interface WeekBucket {
  weekStart: string; // ISO date string (Monday)
  weekEnd: string;
  weekLabel: string; // e.g. "Mar 3"
}

export interface UserWorkload {
  userId: string;
  userName: string;
  userEmail: string;
  userImage: string | null;
  role: string;
  department: string | null;
  maxCapacity: number;
  /** Total active assignments right now (not scoped to week) */
  totalActiveAssignments: number;
  /** Per-week breakdown */
  weeks: WeekData[];
}

export interface WeekData {
  weekStart: string;
  count: number;
  byStatus: Record<string, number>;
  isOverloaded: boolean;
  /** Sum of hours booked via the Resources page grid for this user/week. */
  bookedHours: number;
  /** Weekly capacity cap derived from user.maxCapacity (per-day) × 5 weekdays. */
  weeklyCapacityHours: number;
  assignments: WeekAssignment[];
}

export interface WeekAssignment {
  assignmentId: string;
  stageId: string;
  stageName: string;
  stageStatus: string;
  deliverableName: string;
  projectName: string;
  projectId: string;
  dueDate: string | null;
  priority: string;
}

/**
 * Generate week buckets for a given number of weeks from a start date.
 */
export function generateWeekBuckets(
  numWeeks: number,
  fromDate: Date = new Date()
): WeekBucket[] {
  const buckets: WeekBucket[] = [];
  // Start from the Monday of the current week
  const currentWeekStart = startOfWeek(fromDate, { weekStartsOn: 1 });

  for (let i = 0; i < numWeeks; i++) {
    const ws = addWeeks(currentWeekStart, i);
    const we = endOfWeek(ws, { weekStartsOn: 1 });
    buckets.push({
      weekStart: ws.toISOString(),
      weekEnd: we.toISOString(),
      weekLabel: format(ws, "MMM d"),
    });
  }

  return buckets;
}

/**
 * Get workload data for all users in an organization over a number of weeks.
 */
export async function getWorkloadData(
  organizationId: string,
  ctx: VisibilityContext,
  options: {
    numWeeks?: number;
    projectId?: string;
    stageType?: string;
    userId?: string;
  } = {}
): Promise<{ users: UserWorkload[]; weeks: WeekBucket[] }> {
  const { numWeeks = 8, projectId, stageType, userId } = options;
  const weeks = generateWeekBuckets(numWeeks);

  // Visibility: admins see all; others are scoped to their client teams.
  // Nested assignment queries can't easily be filtered at the DB level, so we
  // pull with project.clientTeamId and filter in application code.
  const isAdmin = ctx.role === "ADMIN";
  const allowedTeamIds = new Set(ctx.clientTeamIds ?? []);

  // Pull ResourceBookings for the full horizon in one shot so we can fold
  // them into the per-week totals below without N+1 queries. Bookings and
  // stage assignments capture different things (ad-hoc hours vs. stage work)
  // so we report them alongside each other rather than double-counting.
  const horizonStart = weeks.length > 0 ? new Date(weeks[0].weekStart) : new Date();
  const horizonEnd =
    weeks.length > 0 ? new Date(weeks[weeks.length - 1].weekEnd) : new Date();
  const bookings = await prisma.resourceBooking.findMany({
    where: {
      organizationId,
      ...(userId ? { userId } : {}),
      date: { gte: horizonStart, lte: horizonEnd },
    },
    select: { userId: true, date: true, hours: true },
  });

  // Aggregate bookings by (userId, weekStartIso) so the inner loop is O(1).
  const bookingHoursByUserWeek = new Map<string, number>();
  for (const b of bookings) {
    const weekStart = startOfWeek(b.date, { weekStartsOn: 1 }).toISOString();
    const key = `${b.userId}_${weekStart}`;
    bookingHoursByUserWeek.set(key, (bookingHoursByUserWeek.get(key) ?? 0) + b.hours);
  }

  // Get all users in the organization with their assignments
  const users = await prisma.user.findMany({
    where: {
      organizationId,
      ...(userId ? { id: userId } : {}),
    },
    select: {
      id: true,
      name: true,
      email: true,
      image: true,
      role: true,
      department: true,
      maxCapacity: true,
      assignments: {
        include: {
          deliverableStage: {
            include: {
              template: true,
              deliverable: {
                include: {
                  project: {
                    select: {
                      id: true,
                      name: true,
                      status: true,
                      organizationId: true,
                      clientTeamId: true,
                    },
                  },
                },
              },
            },
          },
        },
      },
    },
    orderBy: { name: "asc" },
  });

  const result: UserWorkload[] = users.map((user) => {
    // Filter assignments by project/stageType if specified, AND by client-team visibility
    let assignments = user.assignments.filter((a) => {
      const proj = a.deliverableStage.deliverable.project;
      if (proj.organizationId !== organizationId) return false;
      if (projectId && proj.id !== projectId) return false;
      if (stageType && a.deliverableStage.template.slug !== stageType) return false;
      // Visibility: non-admins only see projects on their client teams
      if (!isAdmin) {
        if (!proj.clientTeamId || !allowedTeamIds.has(proj.clientTeamId)) return false;
      }
      return true;
    });

    // Active assignments = those in non-terminal statuses
    const activeAssignments = assignments.filter((a) =>
      (ACTIVE_STATUSES as readonly string[]).includes(a.deliverableStage.status)
    );

    // For each week, determine which assignments are "active" during that week.
    // An assignment is active in a week if:
    //   - The stage is currently in an active status, AND
    //   - The stage doesn't have a completedDate before the week starts, AND
    //   - The stage startDate (or assignment createdAt) is before the week ends
    const weekData: WeekData[] = weeks.map((week) => {
      const weekStartDate = new Date(week.weekStart);
      const weekEndDate = new Date(week.weekEnd);

      const weekAssignments = assignments.filter((a) => {
        const stage = a.deliverableStage;

        // Terminal stages drop out entirely — once a stage is APPROVED
        // / DELIVERED / SKIPPED the assignee is done with it, so it
        // shouldn't keep counting toward their workload even in the
        // week it completed. This was previously showing approved
        // stages in the "just-finished" week which made people look
        // busier than they actually were.
        if ((TERMINAL_STATUSES as readonly string[]).includes(stage.status)) {
          return false;
        }

        // If the stage hasn't started and was created after this week ends, skip
        const stageStartOrCreate = stage.startDate || a.createdAt;
        if (new Date(stageStartOrCreate) > weekEndDate) {
          return false;
        }

        // Stage is currently active — it spans into this week
        return true;
      });

      // Build status breakdown
      const byStatus: Record<string, number> = {};
      for (const a of weekAssignments) {
        const status = a.deliverableStage.status;
        byStatus[status] = (byStatus[status] || 0) + 1;
      }

      const bookedHours =
        bookingHoursByUserWeek.get(`${user.id}_${week.weekStart}`) ?? 0;
      // Weekly capacity = per-day cap × 5 weekdays (Mon–Fri).
      const weeklyCapacityHours = user.maxCapacity * 5;

      return {
        weekStart: week.weekStart,
        count: weekAssignments.length,
        byStatus,
        // Overloaded if EITHER the assignment count exceeds daily capacity OR
        // the booked hours blow past the weekly capacity. Both lenses matter.
        isOverloaded:
          weekAssignments.length > user.maxCapacity ||
          bookedHours > weeklyCapacityHours,
        bookedHours,
        weeklyCapacityHours,
        assignments: weekAssignments.map((a) => ({
          assignmentId: a.id,
          stageId: a.deliverableStage.id,
          stageName: a.deliverableStage.template.name,
          stageStatus: a.deliverableStage.status,
          deliverableName: a.deliverableStage.deliverable.name,
          projectName: a.deliverableStage.deliverable.project.name,
          projectId: a.deliverableStage.deliverable.project.id,
          dueDate: a.deliverableStage.dueDate?.toISOString() ?? null,
          priority: a.deliverableStage.deliverable.priority,
        })),
      };
    });

    return {
      userId: user.id,
      userName: user.name || user.email,
      userEmail: user.email,
      userImage: user.image,
      role: user.role,
      department: user.department,
      maxCapacity: user.maxCapacity,
      totalActiveAssignments: activeAssignments.length,
      weeks: weekData,
    };
  });

  return { users: result, weeks };
}

/**
 * Update a user's max capacity.
 */
export async function updateUserCapacity(userId: string, maxCapacity: number) {
  return prisma.user.update({
    where: { id: userId },
    data: { maxCapacity },
    select: { id: true, name: true, maxCapacity: true },
  });
}