import type { BudgetLine, BudgetScope, BudgetYear } from '@spendlight/shared';
import { and, eq, gte, inArray, lte } from 'drizzle-orm';
import type { AppContext } from '../context';
import { budgets } from '../db/schema';

// Each category has one amount for every month of a year (month 0), and any
// month can have its own amount instead: Gifts at $100 a month, $800 in December.

export function budgetYear(ctx: AppContext, year: number): BudgetYear {
  const rows = ctx.db.select().from(budgets).where(eq(budgets.year, year)).all();
  const lines = new Map<number, BudgetLine>();
  for (const r of rows) {
    const line = lines.get(r.categoryId) ?? { categoryId: r.categoryId, monthly: null, months: Array(12).fill(null) };
    if (r.month === 0) line.monthly = r.amountCents;
    else line.months[r.month - 1] = r.amountCents;
    lines.set(r.categoryId, line);
  }
  return { year, items: [...lines.values()] };
}

export const effective = (line: BudgetLine | undefined, monthIndex: number): number =>
  line ? (line.months[monthIndex] ?? line.monthly ?? 0) : 0;

/** Budgets for each month in the list ("2026-01"...), by category. */
export function budgetMatrix(ctx: AppContext, months: string[]): Map<number, number[]> {
  const years = [...new Set(months.map((m) => Number(m.slice(0, 4))))];
  const byYear = new Map(years.map((y) => [y, new Map(budgetYear(ctx, y).items.map((l) => [l.categoryId, l]))]));
  const out = new Map<number, number[]>();
  months.forEach((month, i) => {
    const lines = byYear.get(Number(month.slice(0, 4)))!;
    for (const [categoryId, line] of lines) {
      const row = out.get(categoryId) ?? Array(months.length).fill(0);
      row[i] = effective(line, Number(month.slice(5, 7)) - 1);
      out.set(categoryId, row);
    }
  });
  return out;
}

function put(ctx: AppContext, categoryId: number, year: number, month: number, amountCents: number, now: string) {
  ctx.db
    .insert(budgets)
    .values({ categoryId, year, month, amountCents, updatedAt: now })
    .onConflictDoUpdate({ target: [budgets.categoryId, budgets.year, budgets.month], set: { amountCents, updatedAt: now } })
    .run();
}

/**
 * Set a category's budget from one month:
 *   month: just that month
 *   rest:  that month through December (earlier months keep what they had)
 *   year:  every month of the year
 */
export function setBudget(ctx: AppContext, categoryId: number, year: number, month: number, amountCents: number, scope: BudgetScope): void {
  const now = new Date().toISOString();
  ctx.sqlite.transaction(() => {
    const before = budgetYear(ctx, year).items.find((l) => l.categoryId === categoryId);
    if (scope === 'month') {
      put(ctx, categoryId, year, month, amountCents, now);
    } else if (scope === 'year') {
      ctx.db.delete(budgets).where(and(eq(budgets.categoryId, categoryId), eq(budgets.year, year))).run();
      if (amountCents > 0) put(ctx, categoryId, year, 0, amountCents, now);
    } else {
      // Earlier months that followed the yearly amount keep their old value.
      for (let m = 1; m < month; m++) {
        if (before?.months[m - 1] == null) {
          const old = before?.monthly ?? 0;
          if (old !== amountCents) put(ctx, categoryId, year, m, old, now);
        }
      }
      ctx.db.delete(budgets).where(and(eq(budgets.categoryId, categoryId), eq(budgets.year, year), gte(budgets.month, month), lte(budgets.month, 12))).run();
      put(ctx, categoryId, year, 0, amountCents, now);
    }
    // Tidy up: overrides equal to the yearly amount add nothing.
    const after = budgetYear(ctx, year).items.find((l) => l.categoryId === categoryId);
    if (after && after.monthly !== null) {
      const same = after.months.flatMap((v, i) => (v === after.monthly ? [i + 1] : []));
      if (same.length) ctx.db.delete(budgets).where(and(eq(budgets.categoryId, categoryId), eq(budgets.year, year), inArray(budgets.month, same))).run();
    }
    if (after && after.monthly === null && after.months.every((v) => v === null || v === 0)) {
      ctx.db.delete(budgets).where(and(eq(budgets.categoryId, categoryId), eq(budgets.year, year))).run();
    }
  })();
}

/** Fill in categories that have no budget for the year yet, as an amount for every month. */
export function fillBudgets(ctx: AppContext, year: number, amounts: Map<number, number>): number {
  const now = new Date().toISOString();
  const existing = new Set(budgetYear(ctx, year).items.map((l) => l.categoryId));
  let filled = 0;
  ctx.sqlite.transaction(() => {
    for (const [categoryId, cents] of amounts) {
      if (existing.has(categoryId) || cents <= 0) continue;
      put(ctx, categoryId, year, 0, cents, now);
      filled++;
    }
  })();
  return filled;
}

/** Copy one year's budgets into another, for categories the new year doesn't have yet. */
export function copyYear(ctx: AppContext, from: number, to: number): number {
  const now = new Date().toISOString();
  const target = new Set(budgetYear(ctx, to).items.map((l) => l.categoryId));
  const rows = ctx.db.select().from(budgets).where(eq(budgets.year, from)).all();
  const copied = new Set<number>();
  ctx.sqlite.transaction(() => {
    for (const r of rows) {
      if (target.has(r.categoryId)) continue;
      put(ctx, r.categoryId, to, r.month, r.amountCents, now);
      copied.add(r.categoryId);
    }
  })();
  return copied.size;
}