echos-ocr/src/lib/ai-ocr.ts

import { generateText, Output } from "ai";
import { gateway } from "@ai-sdk/gateway";
import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
import { z } from "zod";
import { prisma } from "./db";
import type { LanguageModel } from "ai";
import type { FormField as PrismaFormField } from "@/generated/prisma/client";

const DEFAULT_GATEWAY_MODEL = "openai/gpt-4o-mini";
const DEFAULT_OLLAMA_MODEL = "llava:7b";

async function getSettings() {
  let settings = await prisma.appSettings.findUnique({
    where: { id: "singleton" },
  });
  if (!settings) {
    settings = await prisma.appSettings.create({
      data: { id: "singleton" },
    });
  }
  return settings;
}

function getModelInstance(
  provider: string,
  modelId: string,
  ollamaUrl: string
): LanguageModel {
  if (provider === "ollama") {
    const resolvedModel = modelId || DEFAULT_OLLAMA_MODEL;
    const baseURL = (ollamaUrl || process.env.OLLAMA_BASE_URL || "http://192.168.68.108:11434") + "/v1";
    const ollama = createOpenAICompatible({ name: "ollama", baseURL });
    return ollama.chatModel(resolvedModel);
  }

  const resolvedModel = modelId || DEFAULT_GATEWAY_MODEL;
  return gateway(resolvedModel);
}

const responseCardSchema = z.object({
  name: z.string().nullable().describe("Full name as written on the card"),
  gender: z.string().nullable().describe("Male or Female"),
  dateOfBirth: z.string().nullable().describe("Date of birth as written"),
  maritalStatus: z.string().nullable().describe("Married, Single, or Other"),
  maritalStatusOther: z.string().nullable().describe("Value if Other is checked"),
  visitType: z.string().nullable().describe("Which checkbox is marked? Normalize to exactly: 'First/Second Time Guest' or 'Update My Information'. The card may say 'I am a first or second time guest at Echo Life' — map that to 'First/Second Time Guest'. Return null if neither is checked."),
  cellPhone: z.string().nullable().describe("Cell phone number"),
  homePhone: z.string().nullable().describe("Home phone number"),
  email: z.string().nullable().describe("Email address"),
  address: z.string().nullable().describe("Street address"),
  aptNumber: z.string().nullable().describe("Apartment number"),
  city: z.string().nullable().describe("City"),
  state: z.string().nullable().describe("State"),
  zip: z.string().nullable().describe("ZIP code"),
  prayerRequests: z.string().nullable().describe("Written prayer requests"),
  prayerForTeam: z.boolean().describe("Whether prayer team checkbox is checked"),
  prayerConfidential: z.boolean().describe("Whether confidential checkbox is checked"),
  confidence: z.number().min(0).max(100).describe("Your confidence in the OCR accuracy, 0-100"),
});

const surveySchema = z.object({
  messageTopics: z.array(z.string()).describe(
    "ONLY items whose checkbox is physically marked (X, checkmark, or filled). Options: Stress/Anxiety, Marriage, Hearing God's Voice, Dealing With Doubt, Parenting, Grief & Loss, Forgiveness, Finances, Purpose/Calling, Prayer, Healthy Boundaries, Understanding The Bible, Emotional Health, Sharing My Faith, Decision Making, Spiritual Disciplines, Spiritual Gifts. Return empty array if none are marked."
  ),
  messageTopicsOther: z.string().nullable().describe("Value if Other is filled in"),
  nextStep: z.array(z.string()).describe(
    "CRITICAL: This question has EXACTLY 2 separate checkboxes, each on its own line. Inspect each checkbox square individually. 'Baptism' = FIRST checkbox (next to 'Expressing my faith in Jesus...'). 'Next Steps' = SECOND checkbox (next to 'Learning more about becoming a partner...'). A checkbox is checked ONLY if its square contains a visible X, checkmark, or fill. An empty square = NOT checked. Most people check only one. Return empty array if neither is marked."
  ),
  attendanceDuration: z.string().nullable().describe(
    "The single checked radio option. One of: Less than 6 months, 6 Months - 1 Year, 1-3 Years, 4-6 Years, or 7+ Years. null if none is marked."
  ),
  campusPreference: z.array(z.string()).describe(
    "ONLY locations whose checkbox is physically marked. Options: Beulah, Pace/Milton, Gulf Breeze, Warrington. Return empty array if none are marked."
  ),
  campusPreferenceOther: z.string().nullable().describe("Value if Other is filled in"),
  howHeard: z.array(z.string()).describe(
    "ONLY items whose checkbox is physically marked. Options: This is my church home, Regular Attender, Drove by, Social Media, Google, Personal Invite. Return empty array if none are marked."
  ),
  howHeardOther: z.string().nullable().describe("Value if Other is filled in"),
  serviceAttended: z.string().nullable().describe("Service letter: A, B, C, or D"),
  confidence: z.number().min(0).max(100).describe("Your confidence in the OCR accuracy, 0-100"),
});

const RESPONSE_SYSTEM_PROMPT = `You are analyzing a scanned church response card. This is the PERSONAL INFORMATION side.

Extract ALL of the following fields from the image.

CHECKBOX RULES - be very strict:
- A checkbox is CHECKED only if it has a visible mark inside it: an X, a checkmark, a filled square, or pen/pencil marks inside the box.
- An EMPTY box (no marks inside) means UNCHECKED, even if text appears next to it.
- When in doubt, treat a checkbox as UNCHECKED.

For handwritten text, read it as accurately as possible.
Be precise: return null for fields you cannot read. Set prayerForTeam and prayerConfidential to false if the checkboxes are not clearly marked.`;

const SURVEY_SYSTEM_PROMPT = `You are analyzing a scanned church Easter survey form. This is the SURVEY side.

Extract ALL of the following fields from the image.

CHECKBOX RULES - be very strict:
- A checkbox is CHECKED only if its square box contains a visible mark: an X, a checkmark, a filled square, or pen/pencil marks INSIDE the box.
- An EMPTY box (no marks inside the square) means UNCHECKED, even if text appears next to it.
- When in doubt, treat a checkbox as UNCHECKED.
- For array fields (messageTopics, nextStep, campusPreference, howHeard): ONLY include items whose checkbox square is physically marked. Return empty arrays when no checkboxes in that group are marked.
- Do NOT confuse section headings or question titles with checked answers.

SPECIAL ATTENTION for "Next Step" question (question 2): It has exactly 2 checkboxes on separate lines. Look at each checkbox square independently. Most respondents check only one. Do not assume both are checked.

Be precise: return null for fields you cannot read.`;

type FormFieldDef = Pick<
  PrismaFormField,
  "key" | "label" | "type" | "options" | "visibleOnCard" | "visibleOnSurvey"
>;

export function buildSchemaFromFields(fields: FormFieldDef[]) {
  const shape: Record<string, z.ZodTypeAny> = {};
  for (const field of fields) {
    const opts = Array.isArray(field.options)
      ? (field.options as string[])
      : [];
    switch (field.type) {
      case "text":
      case "email":
      case "phone":
      case "textarea":
      case "url":
        shape[field.key] = z.string().nullable().describe(field.label);
        break;
      case "select":
      case "radio":
        shape[field.key] = z
          .string()
          .nullable()
          .describe(
            opts.length > 0
              ? `${field.label}. Options: ${opts.join(", ")}`
              : field.label
          );
        break;
      case "multiselect":
        shape[field.key] = z
          .array(z.string())
          .describe(
            opts.length > 0
              ? `ONLY checked items for ${field.label}. Options: ${opts.join(", ")}`
              : `Checked items for ${field.label}`
          );
        break;
      case "checkbox":
        shape[field.key] = z
          .boolean()
          .describe(`Whether ${field.label} checkbox is checked`);
        break;
      case "date":
        shape[field.key] = z
          .string()
          .nullable()
          .describe(`${field.label} as written`);
        break;
      case "number":
        shape[field.key] = z.number().nullable().describe(field.label);
        break;
      default:
        shape[field.key] = z.string().nullable().describe(field.label);
    }
  }
  shape.confidence = z
    .number()
    .min(0)
    .max(100)
    .describe("Your confidence in the OCR accuracy, 0-100");
  return z.object(shape);
}

export function buildPromptFromFields(
  fields: FormFieldDef[],
  templateName: string
): string {
  const fieldList = fields.map((f) => `- ${f.label} (${f.type})`).join("\n");
  return `You are analyzing a scanned ${templateName}. Extract ALL of the following fields from the image:\n\n${fieldList}\n\nCHECKBOX RULES - be very strict:\n- A checkbox is CHECKED only if it has a visible mark inside it: an X, a checkmark, a filled square, or pen/pencil marks inside the box.\n- An EMPTY box means UNCHECKED.\n- When in doubt, treat a checkbox as UNCHECKED.\n\nFor handwritten text, read it as accurately as possible.\nBe precise: return null for fields you cannot read.`;
}

export interface OcrResult {
  data: Record<string, unknown>;
  confidence: number;
  raw: string;
  side: "response" | "survey";
}

const MAX_RETRIES = 4;
const INITIAL_BACKOFF_MS = 5_000;

async function extractStructured<T extends Record<string, unknown>>(
  model: LanguageModel,
  schema: z.ZodType<T>,
  systemPrompt: string,
  imageBase64: string
): Promise<T> {
  let lastError: unknown;

  for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
    try {
      const { output, text } = await generateText({
        model,
        output: Output.object({ schema }),
        messages: [
          { role: "system", content: systemPrompt },
          {
            role: "user",
            content: [
              { type: "text", text: "Extract all data from this scanned card image." },
              { type: "image", image: Buffer.from(imageBase64, "base64") },
            ],
          },
        ],
      });

      if (!output) {
        throw new Error(`AI model did not return structured output. Raw text: ${(text || "").slice(0, 500)}`);
      }

      return output;
    } catch (err) {
      lastError = err;
      const msg = err instanceof Error ? err.message : String(err);
      const isRateLimit = /rate.?limit|too many requests|429|temporarily|free credits/i.test(msg);

      if (!isRateLimit || attempt === MAX_RETRIES) throw err;

      const backoff = INITIAL_BACKOFF_MS * Math.pow(2, attempt);
      console.log(`[ai-ocr] Rate limited (attempt ${attempt + 1}/${MAX_RETRIES + 1}), retrying in ${backoff / 1000}s...`);
      await new Promise((r) => setTimeout(r, backoff));
    }
  }

  throw lastError;
}

export async function ocrImage(
  imageBase64: string,
  side: "response" | "survey",
  templateFields?: { fields: FormFieldDef[]; templateName: string }
): Promise<OcrResult> {
  const settings = await getSettings();
  const provider = settings.aiProvider || "gateway";
  const modelId = settings.aiModel || "";
  const model = getModelInstance(provider, modelId, settings.ollamaUrl);

  let output: Record<string, unknown>;

  if (templateFields) {
    const relevantFields = templateFields.fields.filter((f) =>
      side === "response" ? f.visibleOnCard : f.visibleOnSurvey
    );
    const dynamicSchema = buildSchemaFromFields(relevantFields);
    const dynamicPrompt = buildPromptFromFields(
      relevantFields,
      templateFields.templateName
    );
    output = await extractStructured(model, dynamicSchema, dynamicPrompt, imageBase64);
  } else if (side === "response") {
    output = await extractStructured(model, responseCardSchema, RESPONSE_SYSTEM_PROMPT, imageBase64);
  } else {
    output = await extractStructured(model, surveySchema, SURVEY_SYSTEM_PROMPT, imageBase64);
  }

  const confidence = typeof output.confidence === "number" ? output.confidence : 50;
  const data = { ...output };
  delete data.confidence;

  return { data, confidence, raw: JSON.stringify(output), side };
}