freezmo/lib/schemas.ts

import { z } from "zod";

/**
 * Zod schemas shared by the API routes (server) and the UI (client).
 * Everything the model returns is validated against these.
 */

export const difficultySchema = z.enum(["easy", "medium", "hard"]);
export type Difficulty = z.infer<typeof difficultySchema>;

/**
 * Photos and rendered pages travel as base64 data URLs — never plain URLs.
 * A URL here would make Ollama fetch an arbitrary address on our behalf (an
 * SSRF hole through the model host), so the shape is part of the contract.
 */
export function dataUrlSchema(maxChars: number) {
  return z
    .string()
    .min(32)
    .max(maxChars)
    .regex(/^data:[a-z0-9][a-z0-9.+-]*\/[a-z0-9][a-z0-9.+-]*;base64,/i, "must be a base64 data URL");
}

const baseActivity = {
  /** Short label shown in the feed chrome, e.g. "Cell biology". */
  topic: z.string().min(1).max(120).optional(),
  explanation: z.string().min(1).max(600),
  /**
   * The words from the student's own notes that carry this answer. The feed
   * shows them back after answering ("From your notes", with the matching
   * passage highlighted), so the quote has to be verbatim.
   */
  source: z.string().min(1).max(1200).optional(),
};

export const flashcardSchema = z.object({
  type: z.literal("flashcard"),
  front: z.string().min(1).max(400),
  back: z.string().min(1).max(1200),
  ...baseActivity,
});

export const multipleChoiceSchema = z.object({
  type: z.literal("multiple_choice"),
  question: z.string().min(1).max(600),
  options: z.array(z.string().min(1).max(300)).length(4),
  correctIndex: z.number().int().min(0).max(3),
  ...baseActivity,
});

export const trueFalseSchema = z.object({
  type: z.literal("true_false"),
  statement: z.string().min(1).max(600),
  answer: z.boolean(),
  ...baseActivity,
});

export const fillBlankSchema = z.object({
  type: z.literal("fill_blank"),
  /** Must contain a `____` placeholder. */
  sentence: z.string().min(1).max(600),
  answer: z.string().min(1).max(200),
  /** Optional word bank; when present it is shown as tappable choices. */
  options: z.array(z.string().min(1).max(200)).min(2).max(6).optional(),
  ...baseActivity,
});

export const matchingSchema = z.object({
  type: z.literal("matching"),
  instruction: z.string().min(1).max(300),
  pairs: z
    .array(
      z.object({
        term: z.string().min(1).max(200),
        definition: z.string().min(1).max(300),
      }),
    )
    .min(3)
    .max(6),
  ...baseActivity,
});

export const orderingSchema = z.object({
  type: z.literal("ordering"),
  instruction: z.string().min(1).max(300),
  /** Items in the CORRECT order; the UI shuffles them. */
  items: z.array(z.string().min(1).max(300)).min(3).max(7),
  ...baseActivity,
});

export const shortAnswerSchema = z.object({
  type: z.literal("short_answer"),
  question: z.string().min(1).max(600),
  referenceAnswer: z.string().min(1).max(1200),
  /** Ideas the grader should look for. */
  keyPoints: z.array(z.string().min(1).max(200)).min(1).max(6),
  ...baseActivity,
});

/**
 * Activities as produced by the model (no id / xp yet - the server adds them).
 */
export const generatedActivitySchema = z.discriminatedUnion("type", [
  flashcardSchema,
  multipleChoiceSchema,
  trueFalseSchema,
  fillBlankSchema,
  matchingSchema,
  orderingSchema,
  shortAnswerSchema,
]);
export type GeneratedActivity = z.infer<typeof generatedActivitySchema>;

/** What the vision step must return. */
export const extractedNotesSchema = z.object({
  title: z.string().min(1).max(120),
  subject: z.string().max(80).optional(),
  summary: z.string().max(400).optional(),
  /**
   * Language of the notes (English name or BCP-47 code). Everything the app
   * generates afterwards must be written in this language.
   */
  language: z.string().min(2).max(60).default("the same language as the notes"),
  /** Clean, structured transcription of everything relevant to study. */
  notes: z.string().min(1),
  keyConcepts: z.array(z.string().min(1).max(300)).min(1).max(40),
  definitions: z
    .array(z.object({ term: z.string().min(1).max(200), definition: z.string().min(1).max(600) }))
    .max(30)
    .default([]),
  formulas: z
    .array(z.object({ name: z.string().min(1).max(200), formula: z.string().min(1).max(400) }))
    .max(30)
    .default([]),
  facts: z.array(z.string().min(1).max(400)).max(40).default([]),
});
export type ExtractedNotes = z.infer<typeof extractedNotesSchema>;

/** Grading response for short answers. */
export const gradeResultSchema = z.object({
  correct: z.boolean(),
  score: z.number().min(0).max(100),
  feedback: z.string().min(1).max(600),
  missing: z.array(z.string().max(200)).max(6).default([]),
});export type GradeResult = z.infer<typeof gradeResultSchema>;

/* ------------------------------------------------------------------ */
/* Make my homework                                                    */
/* ------------------------------------------------------------------ */

export const homeworkModeSchema = z.enum(["make", "guide", "check"]);
export type HomeworkMode = z.infer<typeof homeworkModeSchema>;

export const checkIssueSchema = z.object({
  issue: z.string().min(1).max(600),
  why: z.string().min(1).max(600),
  fix: z.string().max(600).default(""),
});
export type CheckIssue = z.infer<typeof checkIssueSchema>;

/** Mode 3: feedback on the student's own solution. */
export const homeworkCheckSchema = z.object({
  title: z.string().min(1).max(200),
  language: z.string().max(60).optional(),
  verdict: z.string().min(1).max(600),
  correct: z.array(z.string().max(600)).max(12).default([]),
  wrong: z.array(checkIssueSchema).max(12).default([]),
  corrected: z.string().min(1).max(20_000),
});
export type HomeworkCheck = z.infer<typeof homeworkCheckSchema>;

/** Mode 1 — the finished homework, nothing else. */
export const homeworkMakeSchema = z.object({
  title: z.string().min(1).max(200),
  /** English name of the language the assignment was written in. */
  language: z.string().max(60).optional(),
  answer: z.string().min(1).max(20_000),
});
export type HomeworkMake = z.infer<typeof homeworkMakeSchema>;

/**
 * One step of the interactive guide, as the model returns it. `checkpoint` is
 * a normal study activity (validated by the same schema as the feed) that the
 * student answers to check the step landed.
 */
export const homeworkStepSchema = z.object({
  title: z.string().min(1).max(200),
  explanation: z.string().min(1).max(1500),
  /** Progressive help: nudge -> real hint -> the full answer for this step. */
  hints: z.array(z.string().min(1).max(1500)).min(1).max(4),
  /** The full answer for this step; the guide assembles the solution from these. */
  answer: z.string().min(1).max(5000),
  checkpoint: generatedActivitySchema.optional(),
});

export type HomeworkGuide = {
  title: string;
  language?: string;
  steps: z.infer<typeof homeworkStepSchema>[];
};

/* ------------------------------------------------------------------ */
/* Client-side homework types (checkpoints carry id/xp, plus storage)  */
/* ------------------------------------------------------------------ */

export const homeworkCheckpointSchema = generatedActivitySchema.and(
  z.object({ id: z.string(), xp: z.number() }),
);
export type HomeworkCheckpoint = z.infer<typeof homeworkCheckpointSchema>;

export type HomeworkStep = {
  title: string;
  explanation: string;
  hints: string[];
  answer: string;
  checkpoint?: HomeworkCheckpoint;
};

export const homeworkStepStoredSchema: z.ZodType<HomeworkStep> = z.object({
  title: z.string().min(1).max(200),
  explanation: z.string().min(1).max(1500),
  hints: z.array(z.string().min(1).max(1500)).min(1).max(4),
  answer: z.string().min(1).max(5000),
  checkpoint: homeworkCheckpointSchema.optional(),
});

const homeworkBaseFields = {
  id: z.string().min(1).max(80),
  title: z.string().min(1).max(200),
  language: z.string().max(60).optional(),
  createdAt: z.number(),
  updatedAt: z.number(),
  /** The assignment text (photos/documents were reduced to text client-side). */
  sourceText: z.string().max(40_000),
};

export const homeworkRecordSchema = z.discriminatedUnion("mode", [
  z.object({
    ...homeworkBaseFields,
    mode: z.literal("make"),
    answer: z.string().min(1).max(20_000),
  }),
  z.object({
    ...homeworkBaseFields,
    mode: z.literal("guide"),
    steps: z.array(homeworkStepStoredSchema).min(2).max(10),
  }),
  z.object({
    ...homeworkBaseFields,
    mode: z.literal("check"),
    verdict: z.string().min(1).max(600),
    correct: z.array(z.string().max(600)).max(12),
    wrong: z.array(checkIssueSchema).max(12),
    corrected: z.string().min(1).max(20_000),
    /** The student's own attempt, so "Redo" can re-check the same work. */
    attempt: z.string().max(20_000).optional(),
  }),
]);
export type HomeworkRecord = z.infer<typeof homeworkRecordSchema>;
export type HomeworkMakeRecord = Extract<HomeworkRecord, { mode: "make" }>;
export type HomeworkGuideRecord = Extract<HomeworkRecord, { mode: "guide" }>;
export type HomeworkCheckRecord = Extract<HomeworkRecord, { mode: "check" }>;

/* ------------------------------------------------------------------ */
/* Client-side deck types (activity + id/xp, plus storage metadata)     */
/* ------------------------------------------------------------------ */

export type Activity = GeneratedActivity & { id: string; xp: number };

export type Folder = {
  id: string;
  name: string;
  createdAt: number;
};

export const folderSchema: z.ZodType<Folder> = z.object({
  id: z.string().min(1).max(80),
  name: z.string().min(1).max(60),
  createdAt: z.number(),
});

/** Virtual folder that holds every deck without a folder of its own. */
export const UNORGANIZED_ID = "unorganized";

export type Deck = {
  id: string;
  title: string;
  subject?: string;
  summary?: string;
  /** Language the deck is written in (kept for grading, regenerate, add-more). */
  language?: string;
  /** Folder this deck lives in; undefined means "Unorganized". */
  folderId?: string;
  /** Text extracted from the user's notes - reused by regenerate/add-more. */
  sourceText: string;
  difficulty: Difficulty;
  createdAt: number;
  updatedAt: number;
  activities: Activity[];
  attempts: number;
  /** Best result as a 0-100 percentage. */
  bestScore?: number;
};

export const deckSchema: z.ZodType<Deck> = z.object({
  id: z.string(),
  title: z.string(),
  subject: z.string().optional(),
  summary: z.string().optional(),
  language: z.string().optional(),
  folderId: z.string().optional(),
  sourceText: z.string(),
  difficulty: difficultySchema,
  createdAt: z.number(),
  updatedAt: z.number(),
  activities: z.array(generatedActivitySchema.and(z.object({ id: z.string(), xp: z.number() }))).min(1),
  attempts: z.number(),
  bestScore: z.number().optional(),
});

/* ------------------------------------------------------------------ */
/* Learning history (IndexedDB, included in backups and sync)          */
/* ------------------------------------------------------------------ */

export const answerRecordSchema = z.object({
  id: z.string().min(1).max(80),
  deckId: z.string().min(1).max(80),
  deckTitle: z.string().max(200),
  activityId: z.string().min(1).max(80),
  type: z.string().max(40),
  topic: z.string().max(120).optional(),
  correct: z.boolean(),
  at: z.number(),
});
export type AnswerRecord = z.infer<typeof answerRecordSchema>;

export const srsRecordSchema = z.object({
  key: z.string().min(1).max(200),
  deckId: z.string().min(1).max(80),
  activityId: z.string().min(1).max(80),
  ease: z.number().min(1).max(4),
  reps: z.number().int().min(0).max(1000),
  intervalDays: z.number().int().min(0).max(3650),
  lapses: z.number().int().min(0).max(10000),
  dueAt: z.number(),
  lastAt: z.number(),
  lastCorrect: z.boolean(),
});
export type SrsRecord = z.infer<typeof srsRecordSchema>;

/* ------------------------------------------------------------------ */
/* Chat (history, sources, scopes)                                     */
/* ------------------------------------------------------------------ */

export const chatSourceSchema = z.object({
  deckId: z.string().min(1).max(120),
  deckTitle: z.string().max(200),
  section: z.string().max(200),
});
export type ChatSource = z.infer<typeof chatSourceSchema>;

export const chatAttachmentSchema = z.object({
  name: z.string().max(200),
  kind: z.enum(["photo", "document", "text"]),
  chars: z.number().int().min(0).max(10_000_000),
});
export type ChatAttachment = z.infer<typeof chatAttachmentSchema>;

export const chatMessageSchema = z.object({
  id: z.string().min(1).max(80),
  role: z.enum(["user", "assistant"]),
  content: z.string().max(20_000),
  thinking: z.string().max(20_000).optional(),
  model: z.string().max(120).optional(),
  sources: z.array(chatSourceSchema).max(40).optional(),
  attachments: z.array(chatAttachmentSchema).max(10).optional(),
  at: z.number(),
  error: z.object({ message: z.string().max(2000), kind: z.string().max(40).optional() }).optional(),
});
export type ChatMessage = z.infer<typeof chatMessageSchema>;

export const chatScopeSchema = z.discriminatedUnion("kind", [
  z.object({ kind: z.literal("all") }),
  z.object({ kind: z.literal("none") }),
  z.object({ kind: z.literal("decks"), deckIds: z.array(z.string().max(120)).max(200) }),
]);
export type ChatScope = z.infer<typeof chatScopeSchema>;

export const chatConversationSchema = z.object({
  id: z.string().min(1).max(80),
  title: z.string().max(200),
  messages: z.array(chatMessageSchema).max(300),
  scope: chatScopeSchema,
  model: z.string().max(120).optional(),
  provider: z.string().max(60).optional(),
  createdAt: z.number(),
  updatedAt: z.number(),
});
export type ChatConversation = z.infer<typeof chatConversationSchema>;