temetro/backend/src/services/ai/tools.ts

import { tool } from "ai";
import type { UIMessageStreamWriter } from "ai";
import { z } from "zod";

import { eq } from "drizzle-orm";

import { db } from "../../db/index.js";
import { organization } from "../../db/schema/auth.js";
import { appointmentInputSchema } from "../../lib/appointment-validation.js";
import { initialsFromName } from "../../lib/initials.js";
import { inventoryInputSchema } from "../../lib/inventory-validation.js";
import { invoiceInputSchema } from "../../lib/invoice-validation.js";
import { validatePatientImport } from "./import.js";
import { prescriptionInputSchema } from "../../lib/prescription-validation.js";
import { taskInputSchema } from "../../lib/task-validation.js";
import * as analytics from "../analytics.js";
import * as appointments from "../appointments.js";
import * as inventory from "../inventory.js";
import * as patients from "../patients.js";
import * as prescriptions from "../prescriptions.js";
import * as tasks from "../tasks.js";
import type { Patient } from "../../types/patient.js";
import type { Veil } from "./veil.js";

// Context every tool closes over: the caller's clinic + role-derived scoping,
// the Veil safeguard, and the UI stream writer used to push REAL (un-redacted)
// record data to the trusted clinician's screen as custom data parts, while the
// value returned to the model stays Veil-redacted on external providers.
export type ToolContext = {
  orgId: string;
  demographicsOnly: boolean;
  scopeProviderId?: string;
  // The signed-in clinician — needed to scope task visibility (and to stamp the
  // creator when an add is committed via the REST endpoints on approval).
  viewer: { userId: string; userName: string; memberRole: string };
  // The composer's "situation" mode (chat | analysis | graph). In graph mode
  // getPatient renders the record graph instead of the record cards.
  mode?: string;
  veil: Veil;
  writer: UIMessageStreamWriter;
};

// Shared schema for tools that take no real arguments. Google Gemini cannot
// emit a function call when a tool's parameter schema has no properties — it
// prints the call as `tool_code` text instead of invoking it (see the
// previewImport note below). One optional, ignored field keeps the schema
// non-empty so Gemini calls the tool; other providers ignore the extra field.
const emptyToolArgs = z.object({
  filter: z
    .string()
    .optional()
    .describe("Optional free-text filter (ignored — the full list is shown)"),
});

// Compact, model-facing projection of a patient (Veil-redacted upstream). Keeps
// clinical signal, drops bulky arrays the model rarely needs verbatim.
function forModel(p: Patient) {
  return {
    fileNumber: p.fileNumber,
    name: p.name,
    age: p.age,
    sex: p.sex,
    status: p.status,
    pcp: p.pcp,
    allergies: p.allergies,
    alerts: p.alerts,
    problems: p.problems,
    medications: p.medications,
    vitals: p.vitals,
    labs: p.labs,
  };
}

export function createChatTools(ctx: ToolContext) {
  const { orgId, demographicsOnly, scopeProviderId, viewer, mode, veil, writer } =
    ctx;

  // Emit a Chain-of-Thought step to the UI as the agent works. Steps stream live
  // (writer.write flushes immediately) even on the non-streamed external+Veil
  // path, so the clinician always sees progress. Labels must be Veil-safe — use
  // the tokenized identifiers the model passed, never resolved real names.
  let stepSeq = 0;
  function step(label: string) {
    stepSeq += 1;
    writer.write({
      type: "data-step",
      data: { id: `step-${stepSeq}`, label, status: "complete" as const },
    });
  }

  // Register a citable source the model can reference inline. The title is the
  // REAL, clinician-facing label (streamed to the trusted UI, like the cards);
  // the returned id is PHI-free (`s1`, `s2`, …) so it survives Veil rehydration
  // when the model echoes it back as a [[src:id]] marker.
  let sourceSeq = 0;
  function addSource(title: string, kind: string): string {
    sourceSeq += 1;
    const id = `s${sourceSeq}`;
    writer.write({ type: "data-source", data: { id, title, kind } });
    return id;
  }

  // Resolve a possibly-tokenized file number to the real patient record, so an
  // add proposal carries real name/initials into the approval card (the model
  // only ever saw Veil tokens). Returns null when the patient isn't found / is
  // out of scope.
  async function resolvePatient(fileNumber: string): Promise<Patient | null> {
    const real = veil.resolveFileNumber(fileNumber);
    return patients.getPatient(orgId, real, demographicsOnly, scopeProviderId);
  }

  return {
    // Look up one patient by file number (MRN) and show their record cards.
    getPatient: tool({
      description:
        "Retrieve a patient's full record by file number (MRN). Displays it as record cards (or, in Graph mode, as the patient's record graph). Use when the clinician asks about a specific patient.",
      inputSchema: z.object({
        fileNumber: z
          .string()
          .describe("The patient's file number / MRN, e.g. 10293"),
      }),
      execute: async ({ fileNumber }) => {
        step(`Looking up patient ${fileNumber}`);
        const real = veil.resolveFileNumber(fileNumber);
        const patient = await patients.getPatient(
          orgId,
          real,
          demographicsOnly,
          scopeProviderId,
        );
        if (!patient) return { found: false as const, fileNumber };
        // Real data → clinician UI. In graph mode render the record graph; in
        // chat/analysis modes render the record cards. Redacted data → model.
        writer.write(
          mode === "graph"
            ? { type: "data-recordGraph", data: patient }
            : { type: "data-patientCard", data: patient },
        );
        const sourceId = addSource(
          `${patient.name} · MRN ${patient.fileNumber}`,
          "patient",
        );
        return {
          found: true as const,
          sourceId,
          patient: forModel(veil.redactPatient(patient)),
        };
      },
    }),

    // Pull a patient's labs (with high/low flags + trend) and chart them.
    getPatientLabs: tool({
      description:
        "Retrieve a patient's lab results and trend for charting. Use when the clinician asks about labs, results, or values over time.",
      inputSchema: z.object({
        fileNumber: z.string().describe("The patient's file number / MRN"),
      }),
      execute: async ({ fileNumber }) => {
        step(`Reading labs for patient ${fileNumber}`);
        const real = veil.resolveFileNumber(fileNumber);
        const patient = await patients.getPatient(
          orgId,
          real,
          demographicsOnly,
          scopeProviderId,
        );
        if (!patient) return { found: false as const, fileNumber };
        if (demographicsOnly) {
          return { found: false as const, reason: "not_authorized" as const };
        }
        writer.write({
          type: "data-labCard",
          data: {
            fileNumber: patient.fileNumber,
            name: patient.name,
            labs: patient.labs,
            labTrend: patient.labTrend,
          },
        });
        const redacted = veil.redactPatient(patient);
        const sourceId = addSource(`Labs · ${patient.name}`, "lab");
        return {
          found: true as const,
          sourceId,
          name: redacted.name,
          labs: patient.labs,
          labTrend: patient.labTrend,
        };
      },
    }),

    // Search the clinic's patients by name or file number. On a unique match this
    // also displays that patient's record card directly (see below), so the common
    // "show me <name>'s record" flow doesn't depend on a follow-up getPatient call.
    searchPatients: tool({
      description:
        "Search the clinic's patients by name fragment. If EXACTLY ONE patient matches, this already displays their full record card — do NOT call getPatient again; just confirm in one sentence. If multiple match, it returns the matches (with file numbers) so you can disambiguate or call getPatient on the right one.",
      inputSchema: z.object({
        query: z.string().describe("Name or file-number fragment to match"),
      }),
      execute: async ({ query }) => {
        step(`Searching patients for "${query}"`);
        const all = await patients.listPatients(
          orgId,
          demographicsOnly,
          scopeProviderId,
        );
        const q = query.trim().toLowerCase();
        // Keep the full Patient objects so a unique match can render a card.
        const matched = all
          .filter(
            (p) =>
              p.name.toLowerCase().includes(q) ||
              p.fileNumber.toLowerCase().includes(q),
          )
          .slice(0, 10);

        // Unique match → show the record card right here. Gemini often skips the
        // search→getPatient hand-off and just emits a canned sentence, leaving the
        // clinician with no card; rendering it directly removes that dependency.
        if (matched.length === 1) {
          const patient = matched[0]!;
          step(`Looking up patient ${patient.fileNumber}`);
          writer.write(
            mode === "graph"
              ? { type: "data-recordGraph", data: patient }
              : { type: "data-patientCard", data: patient },
          );
          const sourceId = addSource(
            `${patient.name} · MRN ${patient.fileNumber}`,
            "patient",
          );
          return {
            count: 1,
            shownCard: true,
            sourceId,
            patient: forModel(veil.redactPatient(patient)),
          };
        }

        const matches = matched.map((p) => {
          const r = veil.redactPatient(p);
          return { fileNumber: r.fileNumber, name: r.name, status: p.status };
        });
        step(`Found ${matches.length} match(es)`);
        return { count: matches.length, matches };
      },
    }),

    // --- Display the clinic's schedule / work queues (read-only) -------------

    listAppointments: tool({
      description:
        "Display the clinic's appointments. Use when the clinician asks to see the schedule, upcoming visits, or today's appointments.",
      inputSchema: emptyToolArgs,
      execute: async () => {
        step("Loading appointments");
        const all = await appointments.listAppointments(orgId);
        writer.write({ type: "data-appointmentList", data: { appointments: all } });
        // Model-facing rows are Veil-safe: redact patient names to tokens.
        const rows = all.map((a) => ({
          date: a.date,
          time: a.time,
          type: a.type,
          provider: a.provider,
          status: a.status,
          patient: veil.active ? "[PATIENT]" : a.name,
        }));
        const sourceId = addSource("Appointments schedule", "appointments");
        return { count: rows.length, sourceId, appointments: rows };
      },
    }),

    listTasks: tool({
      description:
        "Display the care-team task list. Use when the clinician asks to see open tasks, to-dos, or what's assigned.",
      inputSchema: emptyToolArgs,
      execute: async () => {
        step("Loading tasks");
        const all = await tasks.listTasks(orgId, {
          userId: viewer.userId,
          role: viewer.memberRole,
        });
        writer.write({ type: "data-taskList", data: { tasks: all } });
        const rows = all.map((tk) => ({
          title: tk.title,
          assignee: tk.assignee,
          due: tk.due,
          priority: tk.priority,
          done: tk.done,
        }));
        const sourceId = addSource("Task list", "tasks");
        return { count: rows.length, sourceId, tasks: rows };
      },
    }),

    listPrescriptions: tool({
      description:
        "Display prescriptions for the clinic. Use when the clinician asks to see prescriptions or medications prescribed.",
      inputSchema: emptyToolArgs,
      execute: async () => {
        if (demographicsOnly) {
          return { found: false as const, reason: "not_authorized" as const };
        }
        step("Loading prescriptions");
        const all = await prescriptions.listPrescriptions(orgId);
        writer.write({
          type: "data-prescriptionList",
          data: { prescriptions: all },
        });
        const rows = all.map((rx) => ({
          medication: rx.medication,
          dose: rx.dose,
          frequency: rx.frequency,
          status: rx.status,
          prescribedAt: rx.prescribedAt,
          patient: veil.active ? "[PATIENT]" : rx.name,
        }));
        const sourceId = addSource("Prescriptions", "prescriptions");
        return { count: rows.length, sourceId, prescriptions: rows };
      },
    }),

    // --- Propose an addition (dry-run, NEVER writes) ------------------------
    // Each propose tool validates the record and streams an approval card. The
    // clinician approves in the UI, which commits via the existing RBAC-gated
    // REST endpoint (POST /api/appointments | /tasks | /prescriptions). Nothing
    // is written here.

    proposeAppointment: tool({
      description:
        "Propose a new appointment for the clinician to approve. Does NOT save — it shows an approval card; the clinician confirms before anything is written. Prefer the patient's file number (MRN), which fills name/initials from the record. If the file number is unknown (e.g. parsing a schedule export), pass the patient's name instead; type/provider may be omitted and will be filled with placeholders for the clinician to edit.",
      inputSchema: z.object({
        fileNumber: z
          .string()
          .optional()
          .describe("Patient file number / MRN (may be a token); omit if unknown"),
        name: z
          .string()
          .optional()
          .describe("Patient name — use when no file number is known"),
        date: z.string().describe("Appointment date, YYYY-MM-DD"),
        time: z.string().describe("Appointment time, HH:mm (24h)"),
        type: z
          .string()
          .optional()
          .describe("Visit type, e.g. Follow-up, Consultation"),
        provider: z.string().optional().describe("Provider/clinician name"),
      }),
      execute: async ({ fileNumber, name, date, time, type, provider }) => {
        step(`Drafting appointment for ${fileNumber ?? name ?? "patient"}`);
        const patient = fileNumber ? await resolvePatient(fileNumber) : null;
        // A name (resolved or supplied) is the minimum needed to draft a row.
        const resolvedName = patient?.name ?? (name ? veil.rehydrate(name) : undefined);
        if (!resolvedName) {
          return { ok: false as const, reason: "patient_not_found" as const };
        }
        const candidate = {
          fileNumber: patient?.fileNumber ?? "",
          name: resolvedName,
          initials: patient?.initials ?? "",
          date,
          time,
          type: type ?? "",
          provider: provider ?? "",
          source: "ai" as const,
        };
        const parsed = appointmentInputSchema.safeParse(candidate);
        const issues = parsed.success
          ? []
          : parsed.error.issues.map(
              (i) => `${i.path.join(".") || "(root)"}: ${i.message}`,
            );
        writer.write({
          type: "data-actionPreview",
          data: {
            token: `appt-${stepSeq}`,
            kind: "appointment" as const,
            // Real, ready-to-commit values for the approval card.
            record: parsed.success ? parsed.data : candidate,
            issues,
          },
        });
        return {
          ok: parsed.success,
          issues,
          note: "Preview only — awaiting clinician approval before any write.",
        };
      },
    }),

    proposeTask: tool({
      description:
        "Propose a new care-team task for the clinician to approve. Does NOT save — it shows an approval card the clinician confirms before anything is written.",
      inputSchema: z.object({
        title: z.string().describe("What needs doing"),
        assignee: z.string().optional().describe("Who it's assigned to (free text)"),
        assigneeRole: z
          .enum(["admin", "doctor", "reception", "pharmacy", "lab"])
          .nullish()
          .describe("Department the task belongs to; omit for a personal task"),
        due: z.string().optional().describe("Due date / timeframe (free text)"),
        priority: z.enum(["high", "medium", "low"]).optional(),
        patient: z.string().nullish().describe("Related patient (free text)"),
        notes: z.string().nullish(),
      }),
      execute: async (input) => {
        step(`Drafting task "${input.title}"`);
        // Patient free-text may contain Veil tokens — rehydrate for the card.
        const candidate = {
          ...input,
          patient: input.patient ? veil.rehydrate(input.patient) : input.patient,
        };
        const parsed = taskInputSchema.safeParse(candidate);
        const issues = parsed.success
          ? []
          : parsed.error.issues.map(
              (i) => `${i.path.join(".") || "(root)"}: ${i.message}`,
            );
        writer.write({
          type: "data-actionPreview",
          data: {
            token: `task-${stepSeq}`,
            kind: "task" as const,
            record: parsed.success ? parsed.data : candidate,
            issues,
          },
        });
        return {
          ok: parsed.success,
          issues,
          note: "Preview only — awaiting clinician approval before any write.",
        };
      },
    }),

    proposePrescription: tool({
      description:
        "Propose a new prescription for the clinician to approve. Does NOT save — it shows an approval card the clinician confirms before anything is written. Provide the patient's file number (MRN).",
      inputSchema: z.object({
        fileNumber: z.string().describe("Patient file number / MRN (may be a token)"),
        medication: z.string().describe("Medication name"),
        dose: z.string().optional().describe("Dose, e.g. 500mg"),
        frequency: z.string().describe("Frequency, e.g. twice daily"),
        duration: z.string().nullish().describe("Duration, e.g. 7 days"),
        notes: z.string().nullish(),
      }),
      execute: async ({ fileNumber, medication, dose, frequency, duration, notes }) => {
        if (demographicsOnly) {
          return { ok: false as const, reason: "not_authorized" as const };
        }
        step(`Drafting prescription for patient ${fileNumber}`);
        const patient = await resolvePatient(fileNumber);
        if (!patient) {
          return { ok: false as const, reason: "patient_not_found" as const };
        }
        const candidate = {
          fileNumber: patient.fileNumber,
          name: patient.name,
          initials: patient.initials,
          medication,
          dose: dose ?? "",
          frequency,
          duration: duration ?? null,
          notes: notes ?? null,
          source: "ai" as const,
        };
        const parsed = prescriptionInputSchema.safeParse(candidate);
        const issues = parsed.success
          ? []
          : parsed.error.issues.map(
              (i) => `${i.path.join(".") || "(root)"}: ${i.message}`,
            );
        writer.write({
          type: "data-actionPreview",
          data: {
            token: `rx-${stepSeq}`,
            kind: "prescription" as const,
            record: parsed.success ? parsed.data : candidate,
            issues,
          },
        });
        return {
          ok: parsed.success,
          issues,
          note: "Preview only — awaiting clinician approval before any write.",
        };
      },
    }),

    // --- Clinic-wide reads (aggregates / non-PHI — safe to return to model) ---

    getClinicInfo: tool({
      description:
        "Get the clinic's name and basic info. Use when the clinician asks about their clinic/organization (e.g. 'what's my clinic called?').",
      inputSchema: emptyToolArgs,
      execute: async () => {
        step("Loading clinic info");
        const [org] = await db
          .select({
            name: organization.name,
            slug: organization.slug,
            createdAt: organization.createdAt,
          })
          .from(organization)
          .where(eq(organization.id, orgId));
        const info = {
          name: org?.name ?? "",
          slug: org?.slug ?? null,
          createdAt: org?.createdAt ? org.createdAt.toISOString() : null,
        };
        writer.write({ type: "data-clinicCard", data: info });
        const sourceId = addSource(`Clinic · ${info.name}`, "clinic");
        return { ...info, sourceId };
      },
    }),

    getAnalytics: tool({
      description:
        "Retrieve the clinic's analytics AND earnings — patient/appointment/prescription/task counts plus money billed, paid, and outstanding (from invoices), with a by-month earnings trend. Use for KPIs, earnings, revenue, or performance questions.",
      inputSchema: emptyToolArgs,
      execute: async () => {
        step("Loading clinic analytics");
        const data = await analytics.getAnalytics(orgId);
        writer.write({ type: "data-analyticsCard", data });
        const sourceId = addSource("Clinic analytics", "analytics");
        return { ...data, sourceId }; // aggregates only, no PHI
      },
    }),

    listInventory: tool({
      description:
        "List the clinic's inventory (medications/supplies, stock levels, reorder thresholds). Use for stock, low-stock, or reorder questions.",
      inputSchema: emptyToolArgs,
      execute: async () => {
        step("Loading inventory");
        const items = await inventory.listInventory(orgId);
        writer.write({ type: "data-inventoryList", data: { items } });
        const sourceId = addSource("Inventory", "inventory");
        return {
          count: items.length,
          sourceId,
          items: items.map((i) => ({
            name: i.name,
            form: i.form,
            strength: i.strength,
            stock: i.stockQuantity,
            reorderThreshold: i.reorderThreshold,
          })),
        };
      },
    }),

    proposeInventory: tool({
      description:
        "Propose adding one or more items to the clinic's inventory (medications/supplies) for the clinician to approve — e.g. parse an uploaded stock or purchase list into stock items. Does NOT save; it shows an approval card the clinician confirms. Each item needs a name; form/strength/unit/stockQuantity/reorderThreshold/location/expiresAt (YYYY-MM-DD)/notes are optional. Use this for STOCKING inventory; use proposeInvoice instead when billing a patient for purchased items.",
      inputSchema: z.object({
        items: z
          .array(
            z.object({
              name: z.string().describe("Item / medication name"),
              form: z
                .string()
                .optional()
                .describe("Dosage form, e.g. Tablet, Capsule, Syrup"),
              strength: z.string().optional().describe("e.g. 500mg"),
              unit: z
                .string()
                .optional()
                .describe("Dispensing unit, e.g. box, bottle"),
              stockQuantity: z
                .number()
                .optional()
                .describe("Units currently in stock"),
              reorderThreshold: z
                .number()
                .optional()
                .describe("Low-stock reorder level"),
              location: z.string().optional().describe("Storage location"),
              expiresAt: z
                .string()
                .nullish()
                .describe("Expiry date, YYYY-MM-DD"),
              notes: z.string().nullish(),
            }),
          )
          .describe("Inventory items to add (prices/quantities from the document)"),
      }),
      execute: async ({ items }) => {
        step(`Drafting ${items.length} inventory item(s)`);
        // Inventory is non-PHI — no Veil resolution needed.
        const validated: unknown[] = [];
        const issues: string[] = [];
        items.forEach((item, index) => {
          const parsed = inventoryInputSchema.safeParse(item);
          if (parsed.success) {
            validated.push(parsed.data);
          } else {
            issues.push(
              ...parsed.error.issues.map(
                (i) =>
                  `item ${index + 1} ${i.path.join(".") || "(root)"}: ${i.message}`,
              ),
            );
          }
        });
        writer.write({
          type: "data-actionPreview",
          data: {
            token: `inventory-${stepSeq}`,
            kind: "inventory" as const,
            record: { items: validated.length ? validated : items },
            issues,
          },
        });
        return {
          ok: issues.length === 0,
          count: validated.length,
          issues,
          note: "Preview only — awaiting clinician approval before any write.",
        };
      },
    }),

    proposeInvoice: tool({
      description:
        "Propose a new invoice for the clinician to approve — e.g. parse an uploaded list of purchased medications into billable line items. Does NOT save; it shows an approval card the clinician confirms. Provide the patient/client name (a file number if known) and line items {description, quantity, unitPrice}; prices come from the uploaded document.",
      inputSchema: z.object({
        name: z.string().describe("Patient/client name (may be a token)"),
        fileNumber: z
          .string()
          .optional()
          .describe("Patient file number / MRN if known"),
        lineItems: z
          .array(
            z.object({
              description: z.string(),
              quantity: z.number(),
              unitPrice: z.number(),
            }),
          )
          .describe("Billed items, e.g. each purchased medication"),
        notes: z.string().nullish(),
      }),
      execute: async ({ name, fileNumber, lineItems, notes }) => {
        step(`Drafting invoice for ${fileNumber ?? name}`);
        const patient = fileNumber ? await resolvePatient(fileNumber) : null;
        const resolvedName = patient?.name ?? (name ? veil.rehydrate(name) : "");
        if (!resolvedName) {
          return { ok: false as const, reason: "patient_not_found" as const };
        }
        const candidate = {
          fileNumber: patient?.fileNumber ?? "",
          name: resolvedName,
          initials: patient?.initials ?? initialsFromName(resolvedName),
          lineItems,
          notes: notes ?? null,
          source: "ai" as const,
        };
        const parsed = invoiceInputSchema.safeParse(candidate);
        const issues = parsed.success
          ? []
          : parsed.error.issues.map(
              (i) => `${i.path.join(".") || "(root)"}: ${i.message}`,
            );
        writer.write({
          type: "data-actionPreview",
          data: {
            token: `invoice-${stepSeq}`,
            kind: "invoice" as const,
            record: parsed.success ? parsed.data : candidate,
            issues,
          },
        });
        return {
          ok: parsed.success,
          issues,
          note: "Preview only — awaiting clinician approval before any write.",
        };
      },
    }),

    // Migration: validate parsed records WITHOUT writing. The model parses an
    // uploaded export into our patient shape and calls this; the result drives
    // an approval card. Nothing is inserted until the clinician approves and the
    // client posts to POST /api/ai/import (which re-validates + writes).
    previewImport: tool({
      description:
        "Validate patient records parsed from an uploaded database export, as a dry run. Does NOT save anything. Call this when the clinician wants to import/migrate an existing patient database OR add a single patient; parse the file into our patient shape first. The clinician must approve before any data is written.",
      // A concrete object schema (not z.unknown()): Google Gemini can only emit a
      // real function call when the tool's parameters have a defined JSON schema.
      // An array-of-unknown serializes to an empty schema, which makes Gemini
      // print the call as `tool_code` text instead of invoking it. Validation
      // stays lenient — execute() re-parses each record with patientInputSchema,
      // which coerces gender words, bare-string lists, etc.
      inputSchema: z.object({
        records: z
          .array(
            z.object({
              fileNumber: z
                .string()
                .optional()
                .describe(
                  "File number / MRN; digits only (leave blank to auto-generate)",
                ),
              name: z.string().describe("Patient full name"),
              age: z.number().optional().describe("Age in years"),
              sex: z
                .string()
                .optional()
                .describe("Sex — accepts Male/Female or M/F"),
              status: z
                .string()
                .optional()
                .describe("active, inpatient, or discharged"),
              pcp: z.string().optional().describe("Primary care provider name"),
              alerts: z
                .array(z.string())
                .optional()
                .describe("Free-text clinical alerts"),
              allergies: z
                .array(
                  z.object({
                    substance: z.string().describe("Allergen, e.g. Penicillin"),
                    reaction: z.string().optional(),
                    severity: z
                      .string()
                      .optional()
                      .describe("mild, moderate, or severe"),
                  }),
                )
                .optional(),
              medications: z
                .array(
                  z.object({
                    name: z.string(),
                    dose: z.string().optional(),
                    frequency: z.string().optional(),
                  }),
                )
                .optional(),
              problems: z
                .array(
                  z.object({
                    label: z.string().describe("Problem / diagnosis"),
                    since: z.string().optional(),
                  }),
                )
                .optional(),
              labs: z
                .array(
                  z.object({
                    name: z.string(),
                    value: z.string(),
                    flag: z
                      .string()
                      .optional()
                      .describe("normal, high, low, or critical"),
                    takenAt: z
                      .string()
                      .optional()
                      .describe("Date, YYYY-MM-DD"),
                  }),
                )
                .optional(),
              encounters: z
                .array(
                  z.object({
                    date: z
                      .string()
                      .optional()
                      .describe("Visit date, YYYY-MM-DD"),
                    type: z
                      .string()
                      .optional()
                      .describe("Visit type / department"),
                    provider: z.string().optional(),
                    summary: z
                      .string()
                      .optional()
                      .describe("Diagnosis / treatment / notes combined"),
                  }),
                )
                .optional(),
            }),
          )
          .describe("Patient records parsed from the upload, mapped to temetro's shape"),
      }),
      execute: async ({ records }) => {
        step(`Validating ${records.length} record(s)`);
        const { valid, invalid, total } = validatePatientImport(records);
        // Hand the validated set + the raw records to the UI for an approval
        // card. The client can edit any record, re-validate, and posts the valid
        // set back to /api/ai/import on approval. `records` carries the originals
        // so invalid rows are editable.
        writer.write({
          type: "data-importPreview",
          data: { records, valid, invalid, total },
        });
        step(`${valid.length} ready, ${invalid.length} skipped`);
        return {
          total,
          validCount: valid.length,
          invalidCount: invalid.length,
          invalid,
          note: "Preview only — awaiting clinician approval before any write.",
        };
      },
    }),
  };
}