Momento/memento-note/lib/interactive-page/schema.ts

import { z } from 'zod'
import { INTENT_IDS } from '@/lib/interactive-demo/constants'
import { interactiveDemoV1Schema } from '@/lib/interactive-demo/schema'
import {
  CALLOUT_KINDS,
  INTERACTIVE_PAGE_CAPS,
  INTERACTIVE_PAGE_SCHEMA_VERSION,
  PAGE_BLOCK_TYPES,
  SIM_CAPS,
  STEPS_CAPS,
} from './constants'
import { isValidSimParamId } from './sim-eval'

const intentSchema = z.enum(INTENT_IDS).optional()

const langSchema = z
  .string()
  .regex(/^[a-zA-Z]{2,3}(-[a-zA-Z0-9]{2,8})*$/, 'Invalid BCP-47 language tag')

const chartPayloadSchema = z.object({
  chartType: z.enum(['line', 'bar', 'area']),
  series: z
    .array(
      z.object({
        id: z.string().min(1),
        label: z.string().optional(),
        values: z.array(z.number()),
        intent: intentSchema,
      })
    )
    .min(1),
})

const proseBlock = z.object({
  type: z.literal('prose'),
  md: z.string().min(1),
})

const formulaBlock = z.object({
  type: z.literal('formula'),
  tex: z.string().min(1),
  caption: z.string().optional(),
})

const calloutBlock = z.object({
  type: z.literal('callout'),
  kind: z.enum(CALLOUT_KINDS),
  title: z.string().min(1),
  md: z.string().min(1),
})

const demoBlock = z.object({
  type: z.literal('demo'),
  demo: interactiveDemoV1Schema,
  caption: z.string().optional(),
})

const chartBlock = z.object({
  type: z.literal('chart'),
  payload: chartPayloadSchema,
  caption: z.string().optional(),
})

const statsBlock = z.object({
  type: z.literal('stats'),
  items: z
    .array(
      z.object({
        value: z.string().min(1),
        label: z.string().min(1),
        intent: intentSchema,
      })
    )
    .min(INTERACTIVE_PAGE_CAPS.minStatsItems)
    .max(INTERACTIVE_PAGE_CAPS.maxStatsItems),
})

const tableBlock = z.object({
  type: z.literal('table'),
  columns: z.array(z.string().min(1)).min(1),
  rows: z.array(z.array(z.string())),
  caption: z.string().optional(),
})

const imageBlock = z.object({
  type: z.literal('image'),
  src: z.string().min(1),
  alt: z.string().min(1),
  caption: z.string().optional(),
})

const simParamIdSchema = z
  .string()
  .regex(/^[A-Za-z_][A-Za-z0-9_]*$/, 'Invalid sim identifier')
  .refine(isValidSimParamId, {
    message: 'Identifier collides with a reserved constant/function',
  })

const genericSimParamSchema = z.object({
  id: simParamIdSchema,
  symbol: z.string().min(1),
  label: z.string().min(1),
  min: z.number(),
  max: z.number(),
  step: z.number().positive(),
  defaultValue: z.number(),
  unit: z.string().optional(),
  intent: intentSchema,
})

const genericSimComputedSchema = z.object({
  id: simParamIdSchema,
  symbol: z.string().min(1),
  label: z.string().min(1),
  expr: z.string().min(1).max(SIM_CAPS.maxExprChars),
  unit: z.string().optional(),
  intent: intentSchema,
})

const genericSimSchema = z.object({
  simId: z.literal('generic-formula'),
  title: z.string().min(1),
  intro: z.string().optional(),
  params: z.array(genericSimParamSchema).min(1).max(SIM_CAPS.maxParams),
  computed: z.array(genericSimComputedSchema).min(1).max(SIM_CAPS.maxComputed),
  visual: z.union([
    z.object({ kind: z.literal('gauges') }),
    z.object({ kind: z.literal('bars') }),
    z.object({
      kind: z.literal('curve'),
      xParamId: simParamIdSchema,
      expr: z.string().min(1),
      xLabel: z.string().optional(),
      yLabel: z.string().optional(),
    }),
  ]),
  disclaimer: z.string().optional(),
})

const catalogSimSchema = z.object({
  simId: z
    .string()
    .min(1)
    .refine((s) => s !== 'generic-formula', {
      message: 'generic-formula must use the full inline schema',
    }),
  title: z.string().optional(),
  preset: z.record(z.string(), z.number()).optional(),
  disclaimer: z.string().optional(),
})

const simBlock = z.object({
  type: z.literal('sim'),
  sim: z.union([genericSimSchema, catalogSimSchema]),
  caption: z.string().optional(),
})

const stepsBlock = z.object({
  type: z.literal('steps'),
  title: z.string().optional(),
  steps: z
    .array(
      z.object({
        tex: z.string().min(1),
        rule: z.string().optional(),
        speak: z.string().optional(),
      })
    )
    .min(STEPS_CAPS.minSteps)
    .max(STEPS_CAPS.maxSteps),
  caption: z.string().optional(),
})

export const pageBlockSchema = z.discriminatedUnion('type', [
  proseBlock,
  formulaBlock,
  calloutBlock,
  demoBlock,
  chartBlock,
  statsBlock,
  tableBlock,
  imageBlock,
  simBlock,
  stepsBlock,
])

const sectionSchema = z.object({
  id: z.string().min(1),
  title: z.string().min(1),
  blocks: z
    .array(pageBlockSchema)
    .min(1)
    .max(INTERACTIVE_PAGE_CAPS.maxBlocksPerSection),
})

const overviewSchema = z.object({
  lead: z.string().min(1),
  cards: z
    .array(
      z.object({
        badge: z.string().min(1),
        title: z.string().min(1),
        body: z.string().min(1),
        intent: intentSchema,
      })
    )
    .min(INTERACTIVE_PAGE_CAPS.minOverviewCards)
    .max(INTERACTIVE_PAGE_CAPS.maxOverviewCards),
})

export const pageSpecV1Schema = z.object({
  schemaVersion: z.literal(INTERACTIVE_PAGE_SCHEMA_VERSION),
  id: z.string().min(1),
  lang: langSchema,
  hero: z.object({
    kicker: z.string().min(1),
    title: z.string().min(1),
    subtitle: z.string().optional(),
    meta: z.string().optional(),
  }),
  overview: overviewSchema.optional(),
  sections: z
    .array(sectionSchema)
    .min(2)
    .max(INTERACTIVE_PAGE_CAPS.maxSections),
  footer: z.string().optional(),
})

export type PageSpecV1Parsed = z.infer<typeof pageSpecV1Schema>

/** Re-export for callers. */
export { PAGE_BLOCK_TYPES }