CapaKraken/packages/engine/src/estimate/monthly-spread.ts

/**
 * Monthly spread computation for estimate demand lines.
 *
 * Distributes total hours across months in a date range. Supports:
 * - Even distribution (pro-rated for partial months)
 * - Manual overrides per month
 * - Rebalancing remaining hours after manual edits
 */

/** Format: "YYYY-MM" */
type MonthKey = string;

export interface MonthlySpreadInput {
  totalHours: number;
  startDate: Date;
  endDate: Date;
}

export interface MonthlySpreadResult {
  /** Hours per month, keyed by "YYYY-MM" */
  spread: Record<MonthKey, number>;
  /** Ordered month keys for display */
  months: MonthKey[];
}

export interface RebalanceSpreadInput {
  totalHours: number;
  startDate: Date;
  endDate: Date;
  /** Months with manually locked values */
  lockedMonths: Record<MonthKey, number>;
}

/**
 * Returns ordered month keys between two dates (inclusive).
 */
export function getEstimateMonthRange(startDate: Date, endDate: Date): MonthKey[] {
  const months: MonthKey[] = [];
  const start = new Date(startDate);
  const end = new Date(endDate);

  let cursor = new Date(start.getFullYear(), start.getMonth(), 1);
  const endMonth = new Date(end.getFullYear(), end.getMonth(), 1);

  while (cursor <= endMonth) {
    const year = cursor.getFullYear();
    const month = String(cursor.getMonth() + 1).padStart(2, "0");
    months.push(`${year}-${month}`);
    cursor = new Date(cursor.getFullYear(), cursor.getMonth() + 1, 1);
  }

  return months;
}

/**
 * Counts working days (Mon-Fri) in a date range.
 */
function countWorkingDays(from: Date, to: Date): number {
  let count = 0;
  const cursor = new Date(from);
  cursor.setHours(0, 0, 0, 0);
  const end = new Date(to);
  end.setHours(0, 0, 0, 0);

  while (cursor <= end) {
    const day = cursor.getDay();
    if (day !== 0 && day !== 6) count++;
    cursor.setDate(cursor.getDate() + 1);
  }

  return count;
}

function parseMonthKey(monthKey: string): { year: number; month: number } {
  const parts = monthKey.split("-");
  return {
    year: parseInt(parts[0] ?? "0", 10),
    month: parseInt(parts[1] ?? "1", 10) - 1,
  };
}

/**
 * Get working days per month within a date range, pro-rated for partial months.
 */
function getWorkingDaysPerMonth(
  startDate: Date,
  endDate: Date,
  months: MonthKey[],
): Record<MonthKey, number> {
  const result: Record<MonthKey, number> = {};

  for (const monthKey of months) {
    const { year, month } = parseMonthKey(monthKey);

    const monthStart = new Date(year, month, 1);
    const monthEnd = new Date(year, month + 1, 0); // last day of month

    const effectiveStart = startDate > monthStart ? startDate : monthStart;
    const effectiveEnd = endDate < monthEnd ? endDate : monthEnd;

    result[monthKey] = countWorkingDays(effectiveStart, effectiveEnd);
  }

  return result;
}

/**
 * Distribute total hours evenly across months, weighted by working days.
 */
export function computeEvenSpread(input: MonthlySpreadInput): MonthlySpreadResult {
  const { totalHours, startDate, endDate } = input;
  const months = getEstimateMonthRange(startDate, endDate);

  if (months.length === 0) {
    return { spread: {}, months: [] };
  }

  const workingDays = getWorkingDaysPerMonth(startDate, endDate, months);
  const totalWorkingDays = Object.values(workingDays).reduce((sum, d) => sum + d, 0);

  const spread: Record<MonthKey, number> = {};

  if (totalWorkingDays === 0) {
    // Fallback: distribute evenly by month count
    const perMonth = Math.round((totalHours / months.length) * 10) / 10;
    for (const month of months) {
      spread[month] = perMonth;
    }
  } else {
    for (const month of months) {
      const days = workingDays[month] ?? 0;
      const weight = days / totalWorkingDays;
      spread[month] = Math.round(totalHours * weight * 10) / 10;
    }
  }

  // Adjust rounding error on last month
  const spreadTotal = Object.values(spread).reduce((sum, h) => sum + h, 0);
  const diff = Math.round((totalHours - spreadTotal) * 10) / 10;
  if (diff !== 0 && months.length > 0) {
    const lastMonth = months[months.length - 1]!;
    spread[lastMonth] = Math.round(((spread[lastMonth] ?? 0) + diff) * 10) / 10;
  }

  return { spread, months };
}

/**
 * Rebalance: distribute remaining hours (after locked months) across unlocked months.
 */
export function rebalanceSpread(input: RebalanceSpreadInput): MonthlySpreadResult {
  const { totalHours, startDate, endDate, lockedMonths } = input;
  const months = getEstimateMonthRange(startDate, endDate);

  if (months.length === 0) {
    return { spread: {}, months: [] };
  }

  const lockedTotal = Object.entries(lockedMonths)
    .filter(([key]) => months.includes(key))
    .reduce((sum, [, hours]) => sum + hours, 0);

  const remainingHours = Math.max(0, totalHours - lockedTotal);
  const unlockedMonths = months.filter((m) => !(m in lockedMonths));

  const spread: Record<MonthKey, number> = {};

  // Copy locked months
  for (const month of months) {
    if (month in lockedMonths) {
      spread[month] = lockedMonths[month] ?? 0;
    }
  }

  if (unlockedMonths.length === 0) {
    return { spread, months };
  }

  // Distribute remaining hours across unlocked months
  const workingDays = getWorkingDaysPerMonth(startDate, endDate, unlockedMonths);
  const totalWorkingDays = Object.values(workingDays).reduce((sum, d) => sum + d, 0);

  if (totalWorkingDays === 0) {
    const perMonth = Math.round((remainingHours / unlockedMonths.length) * 10) / 10;
    for (const month of unlockedMonths) {
      spread[month] = perMonth;
    }
  } else {
    for (const month of unlockedMonths) {
      const days = workingDays[month] ?? 0;
      const weight = days / totalWorkingDays;
      spread[month] = Math.round(remainingHours * weight * 10) / 10;
    }
  }

  // Adjust rounding error on last unlocked month (only when locked total doesn't exceed budget)
  if (lockedTotal <= totalHours) {
    const spreadTotal = Object.values(spread).reduce((sum, h) => sum + h, 0);
    const diff = Math.round((totalHours - spreadTotal) * 10) / 10;
    if (diff !== 0 && unlockedMonths.length > 0) {
      const lastUnlocked = unlockedMonths[unlockedMonths.length - 1]!;
      spread[lastUnlocked] = Math.round(((spread[lastUnlocked] ?? 0) + diff) * 10) / 10;
    }
  }

  return { spread, months };
}

/**
 * Summarize monthly spreads across multiple demand lines.
 * Returns total hours per month.
 */
export function summarizeMonthlySpread(
  spreads: Record<MonthKey, number>[],
): Record<MonthKey, number> {
  const totals: Record<MonthKey, number> = {};

  for (const spread of spreads) {
    for (const [month, hours] of Object.entries(spread)) {
      totals[month] = Math.round(((totals[month] ?? 0) + hours) * 10) / 10;
    }
  }

  return totals;
}