HealthLog/src/lib/ai/generate-insight.ts

import type { z } from "zod/v4";
import type { AIProvider, CompletionParams, CompletionResult } from "./types";
import {
  aiInsightResponseSchema,
  findRecommendationsMissingRationale,
  findUncitedRecommendations,
  InsightSchemaError,
  type AIInsightResponse,
  type AIRecommendation,
} from "./schema";
import { computeCitationCoverage } from "./citation-coverage";
import { computeConfidence, type ConfidenceInputs } from "./confidence";
import { annotate } from "@/lib/logging/context";

/**
 * Schema-enforced wrapper around `AIProvider.generateCompletion()`.
 *
 * Phase C1 (v1.4.15): Marc's mandate ("zero hallucinations, must
 * ground on user data") means a free-form LLM response is no longer
 * acceptable. This wrapper:
 *
 *   1. Calls the provider once.
 *   2. Parses the response as JSON.
 *   3. Validates it against `aiInsightResponseSchema` AND cross-checks
 *      that every recommendation's `metricSource` appears in
 *      `citations[]`.
 *   4. On failure: retries ONCE, prepending a corrective system
 *      message that includes the violated zod issues.
 *   5. On second failure: throws `InsightSchemaError` (httpStatus 422).
 *
 * The wrapper does not concern itself with provider-level retries
 * (HTTP 5xx, 401 token-refresh, etc.) — each provider client handles
 * those internally. This wrapper retries only on schema-violation.
 */

export interface GenerateInsightOutcome {
  parsed: AIInsightResponse;
  raw: CompletionResult;
  /** Number of provider calls made (1 on success, 2 on retry). */
  attempts: number;
  /** True iff the retry attempt produced the parsed response. */
  retried: boolean;
}

/**
 * Build the corrective user-prompt suffix that explains to the model
 * what went wrong on the first attempt. Includes the violated zod
 * issues verbatim (truncated to 1KB) so the model can self-correct.
 */
function buildRetryCorrectionMessage(reason: string, details: string): string {
  return `
Your previous response did not satisfy the required JSON schema.
Reason: ${reason}
Details: ${details.slice(0, 1024)}

You MUST return a JSON object matching this schema exactly. Do NOT
include prose, markdown fences, or commentary outside the JSON.

Required top-level fields:
  - summary: string (2-3 sentences)
  - recommendations: array of objects with
      { id: string,
        text: string,
        severity: "info" | "suggestion" | "important" | "urgent",
        metricSource: { type: string, timeRange: string, summary: string, n?: number },
        rationale: { dataWindow: "last7days" | "last30days" | "last90days" | "allTime",
                     comparedTo: non-empty string,
                     deviation: non-empty string } }
  - citations: array of objects with { type: string, timeRange: string, summary: string }
  - warnings: array of objects with { topic: string, message: string, severity?: same enum }

Every recommendation's metricSource MUST also appear in citations[]
(same type + timeRange). Every recommendation MUST carry a
rationale object — dataWindow, comparedTo, and deviation are all
required and rationale.dataWindow MUST equal metricSource.timeRange.
If you cannot ground a recommendation in user data from the
snapshot you were given, OMIT it. An empty recommendations[] is
acceptable.
`;
}

interface TryOutcome {
  ok: boolean;
  parsed?: AIInsightResponse;
  raw: CompletionResult;
  reason?: string;
  details?: string;
  zodIssues?: z.ZodIssue[];
}

async function tryOnce(
  provider: AIProvider,
  params: CompletionParams,
): Promise<TryOutcome> {
  const raw = await provider.generateCompletion(params);

  // Parse JSON.
  let json: unknown;
  try {
    json = JSON.parse(raw.content);
  } catch (e) {
    return {
      ok: false,
      raw,
      reason: "Response was not valid JSON",
      details: (e as Error).message,
    };
  }

  // Schema validate.
  const result = aiInsightResponseSchema.safeParse(json);
  if (!result.success) {
    return {
      ok: false,
      raw,
      reason: "Response did not match the required JSON schema",
      details: result.error.issues
        .slice(0, 8)
        .map((i) => `${i.path.join(".")}: ${i.message}`)
        .join("; "),
      zodIssues: result.error.issues,
    };
  }

  // Cross-check citations.
  const uncited = findUncitedRecommendations(result.data);
  if (uncited.length > 0) {
    const list = uncited
      .slice(0, 4)
      .map(
        (u) =>
          `${u.recommendationId} → ${u.missing.type}/${u.missing.timeRange}`,
      )
      .join(", ");
    return {
      ok: false,
      raw,
      reason: "Recommendations referenced metricSources not in citations[]",
      details: `Missing citations: ${list}`,
    };
  }

  return { ok: true, parsed: result.data, raw };
}

/**
 * Forward the citation-coverage breakdown into the active Wide-Event
 * via `annotate()` (no-op outside a request context). Keys are
 * snake_case to match the rest of the AI Wide-Event vocabulary.
 *
 * v1.4.16 phase B5c — additionally emit `ai_rationale_*` keys so the
 * admin AI quality dashboard can chart rationale-coverage per
 * payload alongside citation-coverage. Rationale is required by the
 * strict schema; the missing-ids list will only ever populate when
 * a legacy payload sneaks through `.passthrough()` — useful as an
 * observability tripwire for the migration window.
 */
function annotateCitationCoverage(parsed: AIInsightResponse): void {
  const coverage = computeCitationCoverage(parsed);
  const rationaleMissing = findRecommendationsMissingRationale(parsed);
  annotate({
    meta: {
      ai_total_recommendations: coverage.totalRecommendations,
      ai_normative_recommendations: coverage.normativeRecommendations,
      ai_cited_normative_recommendations:
        coverage.citedNormativeRecommendations,
      ai_uncited_normative_recommendation_ids:
        coverage.uncitedNormativeRecommendationIds,
      ai_rationale_total_recommendations: coverage.totalRecommendations,
      ai_rationale_with_rationale:
        coverage.totalRecommendations - rationaleMissing.length,
      ai_rationale_missing_recommendation_ids: rationaleMissing,
    },
  });
}

/**
 * v1.4.16 phase B5d — confidence resolver.
 *
 * The wrapper does NOT have direct access to the user's bucket-series
 * data, so the caller can pass an optional resolver that maps each
 * recommendation to its `ConfidenceInputs`. When omitted, the wrapper
 * falls back to inputs derivable from the parsed payload alone:
 *
 *   - n: `metricSource.n` (or 0 when missing — triggers the n<3 hard
 *     cap so a recommendation cited from no quantified data is never
 *     dressed up as confident).
 *   - recencyDays: 0. Defensive default — assumes the cited window
 *     ends "now" because we have no fresher source. The route layer
 *     can supply a real value via the resolver.
 *   - deviationStdRatio: null. Yields the neutral 15-point signal
 *     contribution so a wrapper-only path never over-claims.
 *
 * The resolver fires once per recommendation in payload order so the
 * caller can side-effect (e.g. log per-rec inputs) deterministically.
 */
export type ConfidenceContextResolver = (
  rec: AIRecommendation,
) => ConfidenceInputs;

const DEFAULT_CONFIDENCE_RESOLVER: ConfidenceContextResolver = (rec) => ({
  n: rec.metricSource.n ?? 0,
  recencyDays: 0,
  deviationStdRatio: null,
});

/**
 * Apply the deterministic confidence override to every recommendation.
 * The model's value (if any) is REPLACED. Returns a Wide-Event
 * payload with per-rec model-vs-computed pairs so admin observability
 * can chart drift between what the LLM thinks and what the formula
 * produces.
 */
function applyConfidenceOverride(
  parsed: AIInsightResponse,
  resolver: ConfidenceContextResolver,
): void {
  const overrides: Array<{
    id: string;
    model: number | null;
    computed: number;
  }> = [];
  for (const rec of parsed.recommendations) {
    const inputs = resolver(rec);
    const computed = computeConfidence(inputs);
    const modelValue =
      typeof rec.confidence === "number" ? rec.confidence : null;
    rec.confidence = computed;
    overrides.push({ id: rec.id, model: modelValue, computed });
  }
  if (overrides.length > 0) {
    annotate({
      meta: {
        ai_confidence_override_event: "confidence.override.applied",
        ai_confidence_override_count: overrides.length,
        ai_confidence_override_pairs: overrides,
      },
    });
  }
}

export interface GenerateInsightOptions {
  /**
   * v1.4.16 phase B5d — per-recommendation `ConfidenceInputs`
   * resolver. When supplied, the wrapper uses these inputs to
   * `computeConfidence()` and OVERWRITES `rec.confidence`. When
   * omitted, the wrapper falls back to a parsed-payload-only default
   * (n from `metricSource.n`, recencyDays=0, ratio=null).
   *
   * The override happens regardless — calibrated probabilities are
   * not a small-LLM strength and the deterministic path keeps the
   * v1.4.17 feedback ratchet reproducible.
   */
  confidenceContext?: ConfidenceContextResolver;
}

/**
 * Run a schema-enforced insight generation. Up to 2 provider calls.
 * Throws `InsightSchemaError` (httpStatus 422) if both attempts fail
 * the schema; provider-level errors bubble untouched.
 *
 * v1.4.16 phase B5d — after a successful parse, every recommendation's
 * `confidence` field is overwritten with the deterministic
 * `computeConfidence()` value. The model's number (if any) is
 * discarded; admin observability sees both via the
 * `confidence.override.applied` Wide-Event annotation.
 */
export async function generateInsight(
  provider: AIProvider,
  params: CompletionParams,
  options: GenerateInsightOptions = {},
): Promise<GenerateInsightOutcome> {
  const resolver =
    options.confidenceContext ?? DEFAULT_CONFIDENCE_RESOLVER;
  const first = await tryOnce(provider, params);
  if (first.ok && first.parsed) {
    annotateCitationCoverage(first.parsed);
    applyConfidenceOverride(first.parsed, resolver);
    return {
      parsed: first.parsed,
      raw: first.raw,
      attempts: 1,
      retried: false,
    };
  }

  // Retry-once with corrective context appended to the user prompt.
  const correction = buildRetryCorrectionMessage(
    first.reason ?? "schema mismatch",
    first.details ?? "",
  );
  const retryParams: CompletionParams = {
    ...params,
    userPrompt: `${params.userPrompt}\n\n${correction}`,
  };
  const second = await tryOnce(provider, retryParams);
  if (second.ok && second.parsed) {
    annotateCitationCoverage(second.parsed);
    applyConfidenceOverride(second.parsed, resolver);
    return {
      parsed: second.parsed,
      raw: second.raw,
      attempts: 2,
      retried: true,
    };
  }

  throw new InsightSchemaError(
    second.reason ?? first.reason ?? "Schema validation failed twice",
    {
      issues: second.zodIssues ?? first.zodIssues,
      attempts: 2,
    },
  );
}