proof-of-work/proof-of-work/lib/ai/workoutSchema.ts

import { z } from 'zod';
import { extractJson } from './programSchema';

/**
 * The shape we ask LLMs to produce for a SINGLE day's workout (the
 * "generate today's workout" flow). Distinct from the multi-week
 * AIProgram in programSchema.ts.
 *
 * This does NOT map onto a DB table directly: the user reviews/edits the
 * suggestion, then it pre-populates the normal New Workout form (nothing
 * is persisted until they save through the regular workout path). So the
 * shape is optimized for "pre-fill a logger" not "INSERT a Program".
 *
 * Per exercise we ask for a working `sets` count plus a single target
 * `reps` / `suggestedWeight` — the hand-off expands that into N identical
 * pre-filled SetLogs. (No warmup/ramping distinction in v1.)
 *
 * `exerciseId` is nullable: the model picks from the user's library when
 * it can, but may suggest something not in the library (the preview
 * prompts the user to map it). `exerciseName` is REQUIRED as the display
 * label + fuzzy-match fallback.
 */

export const aiWorkoutExerciseSchema = z.object({
  exerciseId: z.string().nullable(),
  exerciseName: z.string().min(1),
  order: z.number().int().nonnegative(),
  /// Number of working sets to pre-fill. Defaults to 3 in the hand-off
  /// if the model omits it.
  sets: z.number().int().positive().optional().nullable(),
  /// Target reps per set (the user overwrites with what they actually
  /// did). Omit for time/distance-based work.
  reps: z.number().int().positive().optional().nullable(),
  /// Suggested working weight. Null for cardio / bodyweight / stretching.
  suggestedWeight: z.number().nonnegative().optional().nullable(),
  /// "lbs" | "kg". Optional; hand-off falls back to the user's
  /// defaultWeightUnit when null.
  suggestedWeightUnit: z.enum(['lbs', 'kg']).optional().nullable(),
  /// Strength effort (1-10). The hand-off keeps this only for non-cardio
  /// exercises (cardio uses `gear`).
  rpe: z.number().int().min(1).max(10).optional().nullable(),
  /// Cardio breathing gear (1-5). The hand-off keeps this only for
  /// cardio exercises (strength uses `rpe`).
  gear: z.number().int().min(1).max(5).optional().nullable(),
  /// Target duration in seconds for time-based work (e.g. a hold).
  durationSeconds: z.number().int().positive().optional().nullable(),
  notes: z.string().optional().nullable(),
});

export const aiWorkoutSchema = z.object({
  name: z.string().min(1),
  notes: z.string().optional().nullable(),
  exercises: z.array(aiWorkoutExerciseSchema),
});

export type AIWorkout = z.infer<typeof aiWorkoutSchema>;
export type AIWorkoutExercise = z.infer<typeof aiWorkoutExerciseSchema>;

/**
 * JSON-schema-ish doc pasted into the system prompt so the model knows
 * the exact shape to emit (same approach as PROGRAM_OUTPUT_SHAPE — not a
 * provider "structured output" mode, since Ollama support is uneven).
 */
export const WORKOUT_OUTPUT_SHAPE = `{
  "name": "<string, e.g. Upper Body — Shoulder Focus>",
  "notes": "<string, optional, one-line session summary>",
  "exercises": [
    {
      "exerciseId": "<string — REQUIRED — must be an id from the LIBRARY block. If no library exercise fits, pick the closest match and explain in notes; do NOT invent ids.>",
      "exerciseName": "<string, the canonical name from the library>",
      "order": <int >= 0>,
      "sets": <int >= 1, number of working sets>,
      "reps": <int, target reps per set; omit for time/distance work>,
      "suggestedWeight": <number, working weight in the exercise's LIBRARY \`unit\`; omit/null for cardio, bodyweight, stretching>,
      "suggestedWeightUnit": "<\\"lbs\\" | \\"kg\\", optional — match the exercise's \`unit\` from the LIBRARY>",
      "rpe": <int 1-10, strength effort; use for NON-cardio exercises>,
      "gear": <int 1-5, cardio breathing gear; use for CARDIO exercises instead of rpe>,
      "durationSeconds": <int, optional, for timed holds/intervals>,
      "notes": "<string, optional, short coaching cue>"
    }
  ]
}`;

/**
 * Parse + validate a model's raw response into an AIWorkout. Returns a
 * clean workout or a structured error. Mirrors parseAIProgram.
 */
export function parseAIWorkout(
  raw: string,
):
  | { ok: true; workout: AIWorkout }
  | { ok: false; reason: string; json?: string } {
  const json = extractJson(raw);
  if (!json) {
    return {
      ok: false,
      reason: 'Could not find a JSON object in the response.',
    };
  }
  let obj: unknown;
  try {
    obj = JSON.parse(json);
  } catch (e) {
    return {
      ok: false,
      reason: `JSON parse error: ${(e as Error).message}`,
      json,
    };
  }
  const result = aiWorkoutSchema.safeParse(obj);
  if (!result.success) {
    return {
      ok: false,
      reason:
        'JSON did not match the expected shape: ' +
        result.error.errors
          .slice(0, 5)
          .map((e) => `${e.path.join('.')}: ${e.message}`)
          .join('; '),
      json,
    };
  }
  return { ok: true, workout: result.data };
}