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

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),
};

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

export const generatedDeckSchema = z.object({
  title: z.string().min(1).max(120),
  subject: z.string().max(80).optional(),
  summary: z.string().max(400).optional(),
  /** Language of the generated content, e.g. "Czech" or "cs". */
  language: z.string().max(60).optional(),
  activities: z.array(generatedActivitySchema).min(4).max(20),
});
export type GeneratedDeck = z.infer<typeof generatedDeckSchema>;

/** 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"]);
export type HomeworkMode = z.infer<typeof homeworkModeSchema>;

/** 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 const homeworkGuideSchema = z.object({
  title: z.string().min(1).max(200),
  language: z.string().max(60).optional(),
  steps: z.array(homeworkStepSchema).min(2).max(10),
});
export type HomeworkGuide = z.infer<typeof homeworkGuideSchema>;

/* ------------------------------------------------------------------ */
/* 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),
  }),
]);
export type HomeworkRecord = z.infer<typeof homeworkRecordSchema>;
export type HomeworkMakeRecord = Extract<HomeworkRecord, { mode: "make" }>;
export type HomeworkGuideRecord = Extract<HomeworkRecord, { mode: "guide" }>;

/* ------------------------------------------------------------------ */
/* 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(),
});