Jyotisha/frontend/src/lib/consultation-step-timing.ts

/**
 * Per-step clocks for one ordinary consultation.
 *
 * Numbers and statuses only. Provider usage is copied when the provider sent
 * the field; a missing field stays null and is never estimated. Answer text,
 * prompts, and birth data are not fields here.
 */

export type ConsultationModelStepTiming = {
  index: 0 | 1;
  durationMs: number | null;
  reasoningTokens: number | null;
  outputTokens: number | null;
  inputTokens: number | null;
  cachedInputTokens: number | null;
};

export type ConsultationLatencyRecord = {
  steps: Array<ConsultationModelStepTiming | undefined>;
  openIndex: number | null;
  openStartedAt: number | null;
  nextIndex: number;
  step1StartedAt: number | null;
  answerReasoningMs: number | null;
};

export type UsageTimingView = {
  classificationDurationMs: number | null;
  modelSteps: ConsultationModelStepTiming[];
  answerReasoningMs: number | null;
};

type UsageRecord = Record<string, unknown>;

export function createConsultationLatencyRecord(): ConsultationLatencyRecord {
  return {
    steps: [],
    openIndex: null,
    openStartedAt: null,
    nextIndex: 0,
    step1StartedAt: null,
    answerReasoningMs: null,
  };
}

function tokenOrNull(value: unknown): number | null {
  if (typeof value !== "number" || !Number.isFinite(value) || value < 0) return null;
  const truncated = Math.trunc(value);
  return truncated > 1_000_000_000 ? null : truncated;
}

function asRecord(value: unknown): UsageRecord | null {
  return value && typeof value === "object" && !Array.isArray(value)
    ? value as UsageRecord
    : null;
}

/** Read the provider's cached-input count. Absent means null, including when only a sibling field exists. */
export function cachedInputTokensFromUsage(usage: UsageRecord): number | null {
  if ("cachedInputTokens" in usage) return tokenOrNull(usage.cachedInputTokens);
  const details = asRecord(usage.inputTokenDetails);
  if (details && "cacheReadTokens" in details) return tokenOrNull(details.cacheReadTokens);
  const promptDetails = asRecord(usage.prompt_tokens_details);
  if (promptDetails && "cached_tokens" in promptDetails) return tokenOrNull(promptDetails.cached_tokens);
  return null;
}

function usageFromChunk(chunk: { payload?: unknown }): UsageRecord {
  const payload = asRecord(chunk.payload);
  const output = asRecord(payload?.output);
  return asRecord(output?.usage) ?? asRecord(payload?.usage) ?? {};
}

function emptyStep(index: 0 | 1): ConsultationModelStepTiming {
  return {
    index,
    durationMs: null,
    reasoningTokens: null,
    outputTokens: null,
    inputTokens: null,
    cachedInputTokens: null,
  };
}

function closeOpen(record: ConsultationLatencyRecord, now: number, usage: UsageRecord): void {
  const index = record.openIndex;
  const startedAt = record.openStartedAt;
  record.openIndex = null;
  record.openStartedAt = null;
  if (index === null) return;
  record.nextIndex = index + 1;
  if (index !== 0 && index !== 1) return;
  const durationMs = startedAt === null ? null : Math.max(0, Math.trunc(now - startedAt));
  record.steps[index] = {
    index,
    durationMs: durationMs !== null && durationMs > 7 * 24 * 60 * 60 * 1000 ? null : durationMs,
    reasoningTokens: tokenOrNull(usage.reasoningTokens),
    outputTokens: tokenOrNull(usage.outputTokens),
    inputTokens: tokenOrNull(usage.inputTokens),
    cachedInputTokens: cachedInputTokensFromUsage(usage),
  };
}

/** Fold one model-loop chunk into step 0 / step 1. Text payloads are ignored. */
export function observeConsultationModelChunk(
  record: ConsultationLatencyRecord,
  chunk: { type?: string; payload?: unknown },
  now: number,
): void {
  if (chunk.type === "step-start") {
    if (record.openIndex !== null) closeOpen(record, now, {});
    const index = record.nextIndex;
    record.openIndex = index;
    record.openStartedAt = now;
    if (index === 1) record.step1StartedAt = now;
    return;
  }
  if (chunk.type === "step-finish") {
    if (record.openIndex === null) {
      record.openIndex = record.nextIndex;
      record.openStartedAt = null;
    }
    closeOpen(record, now, usageFromChunk(chunk));
  }
}

/**
 * First answer-body character of step 1. Step 0 narration does not count,
 * and a later continuation does not move the clock backward or forward.
 */
export function noteAnswerBodyCharacter(record: ConsultationLatencyRecord, now: number): void {
  if (record.answerReasoningMs !== null) return;
  if (record.openIndex !== 1 || record.step1StartedAt === null) return;
  record.answerReasoningMs = Math.max(0, Math.trunc(now - record.step1StartedAt));
}

export function consultationModelSteps(record: ConsultationLatencyRecord): ConsultationModelStepTiming[] {
  return [emptyStep(0), emptyStep(1)].map((blank, index) => record.steps[index] ?? blank);
}

export function consultationUsageMetadata(input: {
  modelSteps: readonly ConsultationModelStepTiming[] | null | undefined;
  answerReasoningMs: number | null | undefined;
}): { modelSteps: ConsultationModelStepTiming[]; answer: { reasoning_ms: number | null } } {
  const view = consultationUsageTiming({
    classificationDurationMs: null,
    modelSteps: input.modelSteps,
    answerReasoningMs: input.answerReasoningMs,
  });
  return {
    modelSteps: view.modelSteps,
    answer: { reasoning_ms: view.answerReasoningMs },
  };
}

export function consultationUsageTiming(input: {
  classificationDurationMs: number | null;
  modelSteps: readonly ConsultationModelStepTiming[] | null | undefined;
  answerReasoningMs: number | null | undefined;
}): UsageTimingView {
  const steps = input.modelSteps ?? [];
  return {
    classificationDurationMs: input.classificationDurationMs,
    modelSteps: [0, 1].map((index) => {
      const found = steps.find((step) => step.index === index);
      return found ?? emptyStep(index as 0 | 1);
    }),
    answerReasoningMs: input.answerReasoningMs ?? null,
  };
}

function integerOrNull(value: unknown): number | null {
  if (typeof value === "number" && Number.isFinite(value)) return tokenOrNull(value);
  if (typeof value === "string" && value.trim() !== "") {
    const parsed = Number(value);
    return Number.isFinite(parsed) ? tokenOrNull(parsed) : null;
  }
  return null;
}

/** Keep only the numeric timing fields stored on usage_ledger.metadata. */
function parsedJson(value: unknown): unknown {
  if (typeof value !== "string") return value;
  try {
    return JSON.parse(value) as unknown;
  } catch {
    return null;
  }
}

export function usageTimingFromLedger(
  modelSteps: unknown,
  classificationDurationMs: unknown,
  answerReasoningMs: unknown,
): UsageTimingView {
  const decoded = parsedJson(modelSteps);
  const steps = Array.isArray(decoded) ? decoded : [];
  const parsed = steps.map((item) => {
    const record = asRecord(item);
    if (!record || (record.index !== 0 && record.index !== 1)) return null;
    return {
      index: record.index as 0 | 1,
      durationMs: integerOrNull(record.durationMs),
      reasoningTokens: integerOrNull(record.reasoningTokens),
      outputTokens: integerOrNull(record.outputTokens),
      inputTokens: integerOrNull(record.inputTokens),
      cachedInputTokens: integerOrNull(record.cachedInputTokens),
    };
  }).filter((step): step is ConsultationModelStepTiming => step !== null);
  return consultationUsageTiming({
    classificationDurationMs: integerOrNull(classificationDurationMs),
    modelSteps: parsed,
    answerReasoningMs: integerOrNull(answerReasoningMs),
  });
}

function ms(value: number | null): string {
  return value === null ? "—" : `${value} ms`;
}

function tok(value: number | null): string {
  return value === null ? "—" : value.toLocaleString("zh-CN");
}

/** One admin-table cell. Numbers only, so an old row with no timing still reads as dashes. */
export function formatUsageTiming(view: UsageTimingView | null | undefined): string {
  if (!view) return "—";
  const steps = view.modelSteps.length > 0 ? view.modelSteps : [emptyStep(0), emptyStep(1)];
  const lines = [
    `分类 ${ms(view.classificationDurationMs)}`,
    ...steps.map((step) => (
      `第 ${step.index} 步 ${ms(step.durationMs)} · 推理 ${tok(step.reasoningTokens)} · 出 ${tok(step.outputTokens)} · 入 ${tok(step.inputTokens)} · 缓存入 ${tok(step.cachedInputTokens)}`
    )),
    `写到正文 ${ms(view.answerReasoningMs)}`,
  ];
  return lines.join("\n");
}