CapaKraken/packages/api/src/router/report-query-engine.ts

import { TRPCError } from "@trpc/server";
import { z } from "zod";
import { controllerProcedure } from "../trpc.js";
import {
  buildSelect,
  buildWhere,
  csvEscape,
  type EntityKey,
  flattenRow,
  getValidScalarField,
  reportEntitySchema,
  ReportInputSchema,
  type ReportInput,
  type ReportQueryResult,
  validateReportInput,
} from "./report-query-config.js";
import { COLUMN_MAP } from "./report-columns.js";
import { buildResourceMonthReportExplainability } from "./report-explainability.js";
import { buildReportGroups, pickColumns, sortInMemoryRows } from "./report-query-utils.js";
import { executeResourceMonthReport } from "./report-resource-month-query.js";

export const reportQueryProcedures = {
  getAvailableColumns: controllerProcedure
    .input(z.object({ entity: reportEntitySchema }))
    .query(({ input }) => {
      const columns = COLUMN_MAP[input.entity];
      if (!columns) {
        throw new TRPCError({ code: "BAD_REQUEST", message: `Unknown entity: ${input.entity}` });
      }
      return columns.map(({ key, label, dataType }) => ({ key, label, dataType }));
    }),

  getReportData: controllerProcedure
    .input(ReportInputSchema)
    .query(async ({ ctx, input }) => executeReportQuery(ctx.db, input)),

  exportReport: controllerProcedure
    .input(ReportInputSchema.omit({ offset: true }).extend({
      limit: z.number().int().min(1).max(50000).default(5000),
    }))
    .mutation(async ({ ctx, input }) => {
      const result = await executeReportQuery(ctx.db, { ...input, offset: 0 });
      const outputColumns = result.columns;
      const entityColumns = COLUMN_MAP[input.entity];
      const headerLabels = outputColumns.map((key) => {
        const def = entityColumns.find((column) => column.key === key);
        return def?.label ?? key;
      });

      const csvLines = [headerLabels.map(csvEscape).join(",")];
      const groupByLabel = input.groupBy
        ? entityColumns.find((column) => column.key === input.groupBy)?.label ?? input.groupBy
        : null;
      const groupStartByIndex = new Map(
        result.groups.map((group) => [group.startIndex, group] as const),
      );

      result.rows.forEach((row, index) => {
        const group = groupStartByIndex.get(index);
        if (group && groupByLabel) {
          csvLines.push(outputColumns.map((_, columnIndex) => (
            columnIndex === 0 ? csvEscape(`${groupByLabel}: ${group.label} (${group.rowCount})`) : ""
          )).join(","));
        }

        csvLines.push(outputColumns.map((column) => csvEscape(row[column])).join(","));
      });

      return {
        csv: csvLines.join("\n"),
        rowCount: result.rows.length,
        rows: result.rows,
        columns: result.columns,
        groups: result.groups,
        ...(result.explainability ? { explainability: result.explainability } : {}),
      };
    }),
};

async function executeReportQuery(
  db: any,
  input: ReportInput,
): Promise<ReportQueryResult> {
  validateReportInput(input);

  if (input.entity === "resource_month") {
    const result = await executeResourceMonthReport(db, input);
    return {
      ...result,
      explainability: buildResourceMonthReportExplainability(input.columns, input.periodMonth),
    };
  }

  const { entity, columns, filters, sortBy, sortDir, limit, offset } = input;
  const select = buildSelect(entity, columns);
  const where = buildWhere(entity, filters);

  const orderBy: Record<string, string>[] = [];
  if (input.groupBy) {
    const groupField = getValidScalarField(entity, input.groupBy);
    if (!groupField) {
      throw new TRPCError({
        code: "BAD_REQUEST",
        message: `Unsupported group field for ${entity}: ${input.groupBy}`,
      });
    }
    orderBy.push({ [groupField]: "asc" });
  }
  if (sortBy) {
    const validField = getValidScalarField(entity, sortBy);
    if (!validField) {
      throw new TRPCError({
        code: "BAD_REQUEST",
        message: `Unsupported sort field for ${entity}: ${sortBy}`,
      });
    }
    orderBy.push({ [validField]: sortDir });
  }

  const modelDelegate = getModelDelegate(db, entity);
  const [rawRows, totalCount] = await Promise.all([
    modelDelegate.findMany({
      select,
      where,
      ...(orderBy.length > 0 ? { orderBy } : {}),
      take: limit,
      skip: offset,
    }),
    modelDelegate.count({ where }),
  ]);

  const flattenedRows = rawRows.map((row: Record<string, unknown>) => flattenRow(row));
  const rows = sortInMemoryRows(flattenedRows, input.groupBy, sortBy, sortDir, COLUMN_MAP[entity]);
  const outputColumns = ["id", ...columns.filter((column) => column !== "id")];
  const groups = buildReportGroups(rows, input.groupBy);

  return {
    rows: rows.map((row) => pickColumns(row, outputColumns)),
    columns: outputColumns,
    totalCount,
    groups,
  };
}

function getModelDelegate(db: any, entity: EntityKey) {
  switch (entity) {
    case "resource":
      return db.resource;
    case "project":
      return db.project;
    case "assignment":
      return db.assignment;
    default:
      throw new TRPCError({ code: "BAD_REQUEST", message: `Unknown entity: ${entity}` });
  }
}