import { CatalogEntityKind, LabTaskStatus, Prisma } from '@prisma/client';
import { normalizeCatalogLocale } from '../catalog/catalog-label.service';
import {
  prosthesisGroupKey,
  unionWorkflowStepCodes,
} from '../../common/prosthesis-group';

type TransactionClient = Prisma.TransactionClient;

export async function generateLabCaseTasks(
  tx: TransactionClient,
  labCaseId: string,
  localeInput?: string | null,
): Promise<number> {
  const existingCount = await tx.labCaseTask.count({ where: { labCaseId } });
  if (existingCount > 0) {
    return 0;
  }

  const locale = normalizeCatalogLocale(localeInput);

  const toothProsthesisRows = await tx.labCaseToothProsthesis.findMany({
    where: { labCaseId },
    include: {
      detail: { select: { id: true, treatmentType: true, toothSelectionGroups: true } },
      line: { select: { id: true, treatmentType: true, toothSelectionGroups: true } },
    },
  });

  if (toothProsthesisRows.length === 0) {
    return 0;
  }

  const prosthesisCodes = [...new Set(toothProsthesisRows.map((r) => r.prosthesisTypeCode))];

  const prosthesisTypes = await tx.prosthesisType.findMany({
    where: { code: { in: prosthesisCodes } },
    include: {
      steps: {
        orderBy: { stepOrder: 'asc' },
        include: { labWorkflowStep: { select: { code: true } } },
      },
    },
  });

  const stepsByProsthesisCode = new Map(
    prosthesisTypes.map((type) => [
      type.code,
      type.steps.map((s) => ({
        stepOrder: s.stepOrder,
        workflowStepCode: s.labWorkflowStep.code,
      })),
    ]),
  );

  const allStepCodes = [
    ...new Set(
      prosthesisTypes.flatMap((type) =>
        type.steps.map((s) => s.labWorkflowStep.code),
      ),
    ),
  ];

  const stepLabels = await resolveStepLabels(tx, allStepCodes, locale);

  // Per tooth: stacked types become one group key; steps are the union.
  // Teeth are not merged even when they share the same type-set.
  type ToothBucket = {
    sourceKey: string;
    treatmentDetailId: string | null;
    lineId: string | null;
    treatmentType: string;
    selectionGroupId: string;
    codes: Set<string>;
  };

  const byTooth = new Map<string, ToothBucket>();

  for (const row of toothProsthesisRows) {
    const sourceKey =
      row.sourceKey ||
      row.treatmentDetailId ||
      row.lineId ||
      '';
    if (!sourceKey) continue;
    const toothKey = `${sourceKey}::${row.tooth}`;
    const selectionGroupId = row.selectionGroupId?.trim() || '';
    const bucket = byTooth.get(toothKey) ?? {
      sourceKey,
      treatmentDetailId: row.treatmentDetailId,
      lineId: row.lineId,
      treatmentType:
        row.detail?.treatmentType ??
        row.line?.treatmentType ??
        row.treatmentType ??
        'prosthesis',
      selectionGroupId,
      codes: new Set<string>(),
    };
    if (!bucket.selectionGroupId && selectionGroupId) {
      bucket.selectionGroupId = selectionGroupId;
    }
    bucket.codes.add(row.prosthesisTypeCode);
    byTooth.set(toothKey, bucket);
  }

  // One lab item per tooth (or arch sentinel): stacked types on that
  // tooth become one prosthesisTypeCode group key; steps are the union.
  const taskRows: Prisma.LabCaseTaskCreateManyInput[] = [];

  for (const [toothKey, bucket] of byTooth) {
    const tooth = toothKey.slice(bucket.sourceKey.length + 2);
    const codes = [...bucket.codes];
    const lists = codes.map((code) =>
      (stepsByProsthesisCode.get(code) ?? []).map((s) => s.workflowStepCode),
    );
    const unioned = unionWorkflowStepCodes(lists);
    if (unioned.length === 0) {
      continue;
    }

    const groupCode = prosthesisGroupKey(codes);

    unioned.forEach((workflowStepCode, index) => {
      taskRows.push({
        labCaseId,
        treatmentDetailId: bucket.treatmentDetailId,
        lineId: bucket.lineId,
        sourceKey: bucket.sourceKey,
        tooth,
        teeth: [tooth],
        treatmentType: bucket.treatmentType,
        prosthesisTypeCode: groupCode,
        selectionGroupId: bucket.selectionGroupId,
        workflowStepCode,
        stepOrder: index + 1,
        stepLabel: stepLabels.get(workflowStepCode) ?? workflowStepCode,
        status: LabTaskStatus.IN_PROGRESS,
      });
    });
  }

  if (taskRows.length === 0) {
    return 0;
  }

  await tx.labCaseTask.createMany({ data: taskRows });
  return taskRows.length;
}

async function resolveStepLabels(
  tx: TransactionClient,
  stepCodes: string[],
  locale: string,
): Promise<Map<string, string>> {
  if (stepCodes.length === 0) {
    return new Map();
  }

  const rows = await tx.catalogTranslation.findMany({
    where: {
      entityKind: CatalogEntityKind.LAB_WORKFLOW_STEP,
      entityCode: { in: stepCodes },
      locale: { in: [locale, 'en'] },
    },
    select: { entityCode: true, locale: true, label: true },
  });

  const byCode = new Map<string, { en?: string; locale?: string }>();
  for (const row of rows) {
    const entry = byCode.get(row.entityCode) ?? {};
    if (row.locale === 'en') entry.en = row.label;
    if (row.locale === locale) entry.locale = row.label;
    byCode.set(row.entityCode, entry);
  }

  const result = new Map<string, string>();
  for (const code of stepCodes) {
    const entry = byCode.get(code);
    result.set(code, entry?.locale ?? entry?.en ?? code);
  }
  return result;
}