dyolink/backend/src/modules/voice/extraction.wire.ts

import { normalizeFdiCode } from '../../common/fdi';
import type {
  ConnectedSpanIntent,
  DueIntent,
  ProsthesisIntent,
  ToothIntent,
  VoiceIntent,
  Weekday,
} from './voice.types';
import { WEEKDAYS } from './voice.types';

/**
 * The shape the model actually emits, and its JSON schema.
 *
 * Deliberately flat: strict `json_schema` mode has poor support for discriminated unions,
 * so every variant field is present and nullable on the wire. `toVoiceIntent` narrows the
 * flat shape into the internal union the resolvers consume, and is total — anything it
 * cannot classify becomes a shape the resolvers will report as unresolved rather than
 * something that throws here.
 */

export type WireToothIntent = {
  spoken: string;
  /** The two-digit FDI code the clinician spoke; null when the tooth was described. */
  fdi: string | null;
  arch: 'upper' | 'lower' | null;
  side: 'patient_right' | 'patient_left' | null;
  position: number | null;
};

export type WireDue = {
  kind: 'weekday' | 'offset' | 'jalali' | 'gregorian' | 'none';
  weekday: Weekday | null;
  which: 'this' | 'next' | null;
  unit: 'day' | 'week' | 'month' | null;
  amount: number | null;
  jy: number | null;
  jm: number | null;
  jd: number | null;
  y: number | null;
  m: number | null;
  d: number | null;
};

export type WireVoiceIntent = {
  treatmentType: string | null;
  teeth: WireToothIntent[];
  connectedSpans: { from: WireToothIntent; to: WireToothIntent }[];
  comment: string | null;
  prosthesisDefaultType: string | null;
  prosthesisOverrides: { tooth: WireToothIntent; type: string }[];
  labId: string | null;
  labMatchExact: boolean;
  due: WireDue;
};

const TOOTH_SCHEMA = {
  type: 'object',
  additionalProperties: false,
  required: ['spoken', 'fdi', 'arch', 'side', 'position'],
  properties: {
    spoken: {
      type: 'string',
      description: 'The exact transcript words for this tooth.',
    },
    fdi: {
      type: ['string', 'null'],
      description:
        'The two-digit FDI code the clinician said for this tooth, e.g. "26". Null only ' +
        'when the tooth was described in words instead of numbered.',
    },
    arch: { type: ['string', 'null'], enum: ['upper', 'lower', null] },
    side: {
      type: ['string', 'null'],
      enum: ['patient_right', 'patient_left', null],
      description: "The PATIENT's side, never the viewer's.",
    },
    position: {
      type: ['integer', 'null'],
      description:
        'Position from the midline: 1 = central incisor … 8 = third molar. Never an FDI code.',
    },
  },
} as const;

export const VOICE_INTENT_JSON_SCHEMA = {
  type: 'object',
  additionalProperties: false,
  required: [
    'treatmentType',
    'teeth',
    'connectedSpans',
    'comment',
    'prosthesisDefaultType',
    'prosthesisOverrides',
    'labId',
    'labMatchExact',
    'due',
  ],
  properties: {
    treatmentType: {
      type: ['string', 'null'],
      description: 'A treatment type CODE from the supplied list, or null.',
    },
    teeth: { type: 'array', items: TOOTH_SCHEMA },
    connectedSpans: {
      type: 'array',
      description: 'Bridges / splinted units. Endpoints inclusive.',
      items: {
        type: 'object',
        additionalProperties: false,
        required: ['from', 'to'],
        properties: { from: TOOTH_SCHEMA, to: TOOTH_SCHEMA },
      },
    },
    comment: {
      type: ['string', 'null'],
      description: 'Clinical notes, in the spoken language.',
    },
    prosthesisDefaultType: {
      type: ['string', 'null'],
      description:
        'A prosthesis type CODE applied to every tooth unless overridden.',
    },
    prosthesisOverrides: {
      type: 'array',
      items: {
        type: 'object',
        additionalProperties: false,
        required: ['tooth', 'type'],
        properties: { tooth: TOOTH_SCHEMA, type: { type: 'string' } },
      },
    },
    labId: {
      type: ['string', 'null'],
      description: 'An id from the supplied lab list. Never invent one.',
    },
    labMatchExact: {
      type: 'boolean',
      description: 'True only when the spoken name matched a lab name exactly.',
    },
    due: {
      type: 'object',
      additionalProperties: false,
      required: [
        'kind',
        'weekday',
        'which',
        'unit',
        'amount',
        'jy',
        'jm',
        'jd',
        'y',
        'm',
        'd',
      ],
      properties: {
        kind: {
          type: 'string',
          enum: ['weekday', 'offset', 'jalali', 'gregorian', 'none'],
        },
        weekday: { type: ['string', 'null'], enum: [...WEEKDAYS, null] },
        which: { type: ['string', 'null'], enum: ['this', 'next', null] },
        unit: {
          type: ['string', 'null'],
          enum: ['day', 'week', 'month', null],
        },
        amount: { type: ['integer', 'null'] },
        jy: { type: ['integer', 'null'] },
        jm: { type: ['integer', 'null'] },
        jd: { type: ['integer', 'null'] },
        y: { type: ['integer', 'null'] },
        m: { type: ['integer', 'null'] },
        d: { type: ['integer', 'null'] },
      },
    },
  },
} as const;

/** Two digits, quadrant 1-8, position 1-8 — the only thing that can be an FDI code. */
const FDI_SHAPE = /^[1-8][1-8]$/;

function toToothIntent(wire: WireToothIntent | undefined | null): ToothIntent {
  const spoken = typeof wire?.spoken === 'string' ? wire.spoken : '';
  // Persian digits and digit-by-digit dictation ("۲۶", "2 6") are FDI codes that do not
  // match literally; without normalising first they fall through to the positional branch
  // with no quadrant and are reported as unresolved.
  const fdi = normalizeFdiCode(wire?.fdi);
  // Only take the explicit branch for something actually FDI-shaped. A model that emits
  // fdi:"6" alongside correct arch/side/position would otherwise lose the tooth entirely.
  if (FDI_SHAPE.test(fdi)) {
    return { kind: 'explicit', fdi, spoken };
  }
  return {
    kind: 'positional',
    arch: wire?.arch as 'upper' | 'lower',
    side: wire?.side as 'patient_right' | 'patient_left',
    position: typeof wire?.position === 'number' ? wire.position : Number.NaN,
    spoken,
  };
}

function toDueIntent(wire: WireDue | undefined | null): DueIntent | null {
  switch (wire?.kind) {
    case 'weekday':
      return {
        kind: 'weekday',
        weekday: wire.weekday as Weekday,
        which: wire.which as 'this',
      };
    case 'offset':
      return {
        kind: 'offset',
        unit: wire.unit as 'day',
        amount: typeof wire.amount === 'number' ? wire.amount : Number.NaN,
      };
    case 'jalali':
      return {
        kind: 'jalali',
        jy: wire.jy as number,
        jm: wire.jm as number,
        jd: wire.jd as number,
      };
    case 'gregorian':
      return {
        kind: 'gregorian',
        y: wire.y as number,
        m: wire.m as number,
        d: wire.d as number,
      };
    case 'none':
    case undefined:
      return null;
    default:
      // An unrecognised kind means a deadline WAS spoken and we failed to classify it.
      // Passing it through lets the resolver flag it; collapsing it to null would make a
      // misunderstood deadline indistinguishable from no deadline at all.
      return { kind: wire?.kind } as unknown as DueIntent;
  }
}

export function toVoiceIntent(wire: WireVoiceIntent): VoiceIntent {
  const teeth = Array.isArray(wire?.teeth) ? wire.teeth : [];
  const spans = Array.isArray(wire?.connectedSpans) ? wire.connectedSpans : [];
  const overrides = Array.isArray(wire?.prosthesisOverrides)
    ? wire.prosthesisOverrides
    : [];

  const connectedSpans: ConnectedSpanIntent[] = spans.map((span) => ({
    from: toToothIntent(span?.from),
    to: toToothIntent(span?.to),
  }));

  const hasProsthesis =
    wire?.prosthesisDefaultType != null || overrides.length > 0;
  const prosthesis: ProsthesisIntent | null = hasProsthesis
    ? {
        defaultType: wire?.prosthesisDefaultType ?? null,
        overrides: overrides.map((o) => ({
          tooth: toToothIntent(o?.tooth),
          type: o?.type,
        })),
      }
    : null;

  return {
    treatmentType: wire?.treatmentType ?? null,
    teeth: teeth.map(toToothIntent),
    connectedSpans,
    comment: wire?.comment ?? null,
    prosthesis,
    labId: wire?.labId ?? null,
    labMatchExact: wire?.labMatchExact === true,
    due: toDueIntent(wire?.due),
  };
}