import { z } from 'zod';

/**
 * The shape we ask LLMs to produce, validated server-side via Zod
 * after parsing whatever JSON came back. Maps 1:1 onto the existing
 * Program -> ProgramWeek -> ProgramDay -> ProgramExercise tables
 * so the apply step is just a transactional INSERT.
 *
 * `exerciseId` is nullable: the model picks from the user's library
 * when it can but is allowed to suggest an exercise that doesn't
 * exist yet (we'll prompt the user to create it during preview).
 * `exerciseName` is REQUIRED so we always have a display label and a
 * fallback for fuzzy matching.
 */

export const aiExerciseSchema = z.object({
  exerciseId: z.string().nullable(),
  exerciseName: z.string().min(1),
  order: z.number().int().nonnegative(),
  sets: z.number().int().positive().optional().nullable(),
  repsMin: z.number().int().positive().optional().nullable(),
  repsMax: z.number().int().positive().optional().nullable(),
  rpe: z.number().int().min(1).max(10).optional().nullable(),
  restSeconds: z.number().int().nonnegative().optional().nullable(),
  /// Suggested starting weight. Not required (cardio, bodyweight,
  /// stretching all leave it null). When provided alongside an
  /// exerciseId that the user starts a workout from, this seeds the
  /// SetLog.weight as a target.
  suggestedWeight: z.number().nonnegative().optional().nullable(),
  /// "lbs" | "kg". Optional; the apply step falls back to the user's
  /// `defaultWeightUnit` preference when null.
  suggestedWeightUnit: z.enum(['lbs', 'kg']).optional().nullable(),
  notes: z.string().optional().nullable(),
});

export const aiDaySchema = z.object({
  dayOfWeek: z.number().int().min(0).max(6),
  name: z.string().optional().nullable(),
  description: z.string().optional().nullable(),
  exercises: z.array(aiExerciseSchema),
});

export const aiWeekSchema = z.object({
  weekNumber: z.number().int().positive(),
  phase: z.string().optional().nullable(),
  description: z.string().optional().nullable(),
  days: z.array(aiDaySchema),
});

export const aiProgramSchema = z.object({
  name: z.string().min(1),
  description: z.string().optional().nullable(),
  type: z.string().min(1),
  durationWeeks: z.number().int().positive(),
  weeks: z.array(aiWeekSchema),
});

export type AIProgram = z.infer<typeof aiProgramSchema>;
export type AIWeek = z.infer<typeof aiWeekSchema>;
export type AIDay = z.infer<typeof aiDaySchema>;
export type AIExercise = z.infer<typeof aiExerciseSchema>;

/**
 * The JSON-schema-ish description we paste into the system prompt so
 * the model knows the exact shape to emit. We don't pass it to a
 * provider's "structured output" mode (Ollama in particular doesn't
 * support that uniformly across models) — it's just a doc the model
 * reads.
 */
export const PROGRAM_OUTPUT_SHAPE = `{
  "name": "<string>",
  "description": "<string, optional>",
  "type": "<string: hypertrophy | strength | power | endurance | recovery | general>",
  "durationWeeks": <int >= 1>,
  "weeks": [
    {
      "weekNumber": <int >= 1>,
      "phase": "<string, optional, e.g. Volume / Intensity / Deload>",
      "description": "<string, optional>",
      "days": [
        {
          "dayOfWeek": <int 0-6, 0=Sunday>,
          "name": "<string, optional, e.g. Push Day>",
          "description": "<string, optional>",
          "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, optional but recommended>,
              "repsMin": <int, optional>,
              "repsMax": <int, optional>,
              "rpe": <int 1-10, optional>,
              "restSeconds": <int >= 0, optional>,
              "suggestedWeight": <number, optional — starting weight; omit/null for cardio, bodyweight, stretching>,
              "suggestedWeightUnit": "<\\"lbs\\" | \\"kg\\", optional — defaults to user's preferred unit>",
              "notes": "<string, optional, coaching note>"
            }
          ]
        }
      ]
    }
  ]
}`;

/**
 * Try to extract a JSON object from a model's raw response. Models
 * sometimes wrap output in ```json fences or add commentary before/
 * after. This pulls the first balanced {...} block.
 */
export function extractJson(raw: string): string | null {
  // Strip code fences if present
  const fenced = raw.match(/```(?:json)?\s*([\s\S]+?)\s*```/);
  if (fenced) return fenced[1].trim();

  // Otherwise find the first balanced {...}
  const start = raw.indexOf('{');
  if (start < 0) return null;
  let depth = 0;
  let inStr = false;
  let escape = false;
  for (let i = start; i < raw.length; i++) {
    const c = raw[i];
    if (escape) {
      escape = false;
      continue;
    }
    if (c === '\\') {
      escape = true;
      continue;
    }
    if (c === '"') {
      inStr = !inStr;
      continue;
    }
    if (inStr) continue;
    if (c === '{') depth++;
    else if (c === '}') {
      depth--;
      if (depth === 0) return raw.slice(start, i + 1);
    }
  }
  return null;
}

/**
 * Parse + validate a model's raw response. Returns either a clean
 * AIProgram or a structured error.
 */
export function parseAIProgram(
  raw: string,
):
  | { ok: true; program: AIProgram }
  | { 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 = aiProgramSchema.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, program: result.data };
}
