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";

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(
    "This question has exactly 2 checkboxes. Include ONLY the ones with a visible mark (X, checkmark, filled). 'Baptism' = the checkbox for 'Expressing my faith in Jesus / baptized'. 'Next Steps' = the checkbox for 'Learning more about becoming a partner / attend Next Steps'. Return empty array if neither checkbox 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 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 array fields (messageTopics, nextStep, campusPreference, howHeard): ONLY include items whose checkbox 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.

Be 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"
): 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 (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 };
}