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

const PARTIAL_DENTURE_CODE = 'partial_denture';

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.
  // Unconnected teeth stay separate even when they share a type.
  // Connected (bridge) teeth that share a type merge into one pipeline; an extra
  // type on one unit of that bridge is its own pipeline. partial_denture is one
  // lab item per sourceKey.
  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);
  }

  const teethByBridge = new Map<string, Set<string>>();
  for (const [toothKey, bucket] of byTooth) {
    if (!bucket.selectionGroupId) continue;
    const tooth = toothKey.slice(bucket.sourceKey.length + 2);
    const bridgeKey = `${bucket.sourceKey}::${bucket.selectionGroupId}`;
    const set = teethByBridge.get(bridgeKey) ?? new Set<string>();
    set.add(tooth);
    teethByBridge.set(bridgeKey, set);
  }
  const bridgeKeys = new Set(
    [...teethByBridge.entries()]
      .filter(([, teeth]) => teeth.size > 1)
      .map(([key]) => key),
  );

  type Pipeline = {
    sourceKey: string;
    treatmentDetailId: string | null;
    lineId: string | null;
    treatmentType: string;
    selectionGroupId: string;
    tooth: string;
    teeth: string[];
    codes: string[];
  };

  type PartialMerge = {
    sourceKey: string;
    treatmentDetailId: string | null;
    lineId: string | null;
    treatmentType: string;
    teeth: string[];
    groupIds: Set<string>;
  };

  type ConnectedTypeMerge = {
    sourceKey: string;
    treatmentDetailId: string | null;
    lineId: string | null;
    treatmentType: string;
    selectionGroupId: string;
    code: string;
    teeth: string[];
  };

  const pipelines: Pipeline[] = [];
  const partialMerges = new Map<string, PartialMerge>();
  const connectedTypeMerges = new Map<string, ConnectedTypeMerge>();

  for (const [toothKey, bucket] of byTooth) {
    const tooth = toothKey.slice(bucket.sourceKey.length + 2);
    const codes = new Set(bucket.codes);
    if (codes.has(PARTIAL_DENTURE_CODE)) {
      codes.delete(PARTIAL_DENTURE_CODE);
      const merge = partialMerges.get(bucket.sourceKey) ?? {
        sourceKey: bucket.sourceKey,
        treatmentDetailId: bucket.treatmentDetailId,
        lineId: bucket.lineId,
        treatmentType: bucket.treatmentType,
        teeth: [],
        groupIds: new Set<string>(),
      };
      merge.teeth.push(tooth);
      if (bucket.selectionGroupId) merge.groupIds.add(bucket.selectionGroupId);
      partialMerges.set(bucket.sourceKey, merge);
    }
    if (codes.size === 0) continue;

    const bridgeKey = bucket.selectionGroupId
      ? `${bucket.sourceKey}::${bucket.selectionGroupId}`
      : '';
    if (bridgeKey && bridgeKeys.has(bridgeKey)) {
      for (const code of codes) {
        const typeKey = `${bridgeKey}::${code}`;
        const merge = connectedTypeMerges.get(typeKey) ?? {
          sourceKey: bucket.sourceKey,
          treatmentDetailId: bucket.treatmentDetailId,
          lineId: bucket.lineId,
          treatmentType: bucket.treatmentType,
          selectionGroupId: bucket.selectionGroupId,
          code,
          teeth: [],
        };
        if (!merge.teeth.includes(tooth)) merge.teeth.push(tooth);
        connectedTypeMerges.set(typeKey, merge);
      }
      continue;
    }

    pipelines.push({
      sourceKey: bucket.sourceKey,
      treatmentDetailId: bucket.treatmentDetailId,
      lineId: bucket.lineId,
      treatmentType: bucket.treatmentType,
      selectionGroupId: bucket.selectionGroupId,
      tooth,
      teeth: [tooth],
      codes: [...codes],
    });
  }

  for (const merge of connectedTypeMerges.values()) {
    const teeth = sortProsthesisTeeth(merge.teeth);
    pipelines.push({
      sourceKey: merge.sourceKey,
      treatmentDetailId: merge.treatmentDetailId,
      lineId: merge.lineId,
      treatmentType: merge.treatmentType,
      selectionGroupId: merge.selectionGroupId,
      tooth: teeth[0] ?? '',
      teeth,
      codes: [merge.code],
    });
  }

  for (const merge of partialMerges.values()) {
    const teeth = sortProsthesisTeeth(merge.teeth);
    const sharedGroup = merge.groupIds.size === 1 ? [...merge.groupIds][0] : '';
    pipelines.push({
      sourceKey: merge.sourceKey,
      treatmentDetailId: merge.treatmentDetailId,
      lineId: merge.lineId,
      treatmentType: merge.treatmentType,
      selectionGroupId: sharedGroup,
      tooth: teeth[0] ?? '',
      teeth,
      codes: [PARTIAL_DENTURE_CODE],
    });
  }

  const taskRows: Prisma.LabCaseTaskCreateManyInput[] = [];

  for (const pipeline of pipelines) {
    const lists = pipeline.codes.map((code) =>
      (stepsByProsthesisCode.get(code) ?? []).map((s) => s.workflowStepCode),
    );
    const unioned = unionWorkflowStepCodes(lists);
    if (unioned.length === 0) {
      continue;
    }

    const groupCode = prosthesisGroupKey(pipeline.codes);

    unioned.forEach((workflowStepCode, index) => {
      taskRows.push({
        labCaseId,
        treatmentDetailId: pipeline.treatmentDetailId,
        lineId: pipeline.lineId,
        sourceKey: pipeline.sourceKey,
        tooth: pipeline.tooth,
        teeth: pipeline.teeth,
        treatmentType: pipeline.treatmentType,
        prosthesisTypeCode: groupCode,
        selectionGroupId: pipeline.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;
}