import type { InferenceUsageInput } from "./inference-accounting";

/**
 * Usage extraction for the three passthrough dialects. The controller never
 * rewrites what an engine says — these helpers only read token counts out of
 * response payloads (and SSE frames) so requests can be recorded.
 *
 * Field spellings differ per dialect: chat completions reports
 * prompt_tokens/completion_tokens, the Responses API and the Anthropic
 * Messages API both report input_tokens/output_tokens with their own cache and
 * reasoning detail shapes. Everything funnels into InferenceUsageInput.
 */
export type ProxyDialect = "chat" | "responses" | "messages";

const asRecord = (value: unknown): Record<string, unknown> | null =>
  value && typeof value === "object" && !Array.isArray(value)
    ? (value as Record<string, unknown>)
    : null;

const asCount = (value: unknown): number | undefined =>
  typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined;

type CountField = "prompt_tokens" | "completion_tokens" | "reasoning_tokens" | "cache_read_tokens" | "cache_write_tokens";

const setCount = (target: InferenceUsageInput, field: CountField, value: unknown): void => {
  const count = asCount(value);
  if (count !== undefined) target[field] = count;
};

const chatUsage = (usage: Record<string, unknown>): InferenceUsageInput => {
  const result: InferenceUsageInput = {};
  setCount(result, "prompt_tokens", usage["prompt_tokens"]);
  setCount(result, "completion_tokens", usage["completion_tokens"]);
  setCount(result, "reasoning_tokens", usage["reasoning_tokens"]);
  const promptDetails = asRecord(usage["prompt_tokens_details"]);
  const completionDetails = asRecord(usage["completion_tokens_details"]);
  if (promptDetails) result.prompt_tokens_details = promptDetails as Record<string, number>;
  if (completionDetails) {
    result.completion_tokens_details = completionDetails as Record<string, number>;
  }
  return result;
};

const responsesUsage = (usage: Record<string, unknown>): InferenceUsageInput => {
  const result: InferenceUsageInput = {};
  setCount(result, "prompt_tokens", usage["input_tokens"]);
  setCount(result, "completion_tokens", usage["output_tokens"]);
  setCount(result, "reasoning_tokens", asRecord(usage["output_tokens_details"])?.["reasoning_tokens"]);
  setCount(result, "cache_read_tokens", asRecord(usage["input_tokens_details"])?.["cached_tokens"]);
  return result;
};

const messagesUsage = (usage: Record<string, unknown>): InferenceUsageInput => {
  const result: InferenceUsageInput = {};
  setCount(result, "prompt_tokens", usage["input_tokens"]);
  setCount(result, "completion_tokens", usage["output_tokens"]);
  setCount(result, "cache_read_tokens", usage["cache_read_input_tokens"]);
  setCount(result, "cache_write_tokens", usage["cache_creation_input_tokens"]);
  return result;
};

/** Usage out of one payload: a non-streaming body, or one parsed SSE data frame. */
export const usageFromPayload = (
  dialect: ProxyDialect,
  payload: Record<string, unknown>,
): InferenceUsageInput | null => {
  // Streaming envelopes nest the usage: Responses under `response` on
  // response.completed, Messages under `message` on message_start.
  const envelope =
    dialect === "responses"
      ? (asRecord(payload["response"]) ?? payload)
      : dialect === "messages"
        ? (asRecord(payload["message"]) ?? payload)
        : payload;
  const usage = asRecord(envelope["usage"]);
  if (!usage) return null;
  const extracted =
    dialect === "chat"
      ? chatUsage(usage)
      : dialect === "responses"
        ? responsesUsage(usage)
        : messagesUsage(usage);
  return Object.keys(extracted).length > 0 ? extracted : null;
};

export interface UsageObserverCallbacks {
  /** Called with every usage payload seen; later fields overwrite earlier ones
   *  (the Messages dialect splits input and output across two events). */
  onUsage: (usage: InferenceUsageInput) => void;
  onFirstFrame: () => void;
}

/**
 * A pure observer for SSE bodies: every chunk is forwarded byte-for-byte and
 * unmodified, while a parallel line buffer splits frames to read usage. This is
 * what "pass through and deal with recording" means mechanically — the client
 * sees exactly what the engine sent.
 */
export const createUsageObserver = (
  dialect: ProxyDialect,
  callbacks: UsageObserverCallbacks,
): TransformStream<Uint8Array, Uint8Array> => {
  const decoder = new TextDecoder();
  let buffer = "";
  let sawFrame = false;

  const observeLine = (line: string): void => {
    const trimmed = line.trim();
    if (!trimmed.startsWith("data:")) return;
    if (!sawFrame) {
      sawFrame = true;
      callbacks.onFirstFrame();
    }
    const data = trimmed.slice("data:".length).trim();
    if (!data || data === "[DONE]") return;
    try {
      const parsed = JSON.parse(data) as unknown;
      const record = asRecord(parsed);
      if (!record) return;
      const usage = usageFromPayload(dialect, record);
      if (usage) callbacks.onUsage(usage);
    } catch {
      // Partial or non-JSON frame: recording is best-effort, forwarding is not.
    }
  };

  return new TransformStream<Uint8Array, Uint8Array>({
    transform(chunk, controller): void {
      controller.enqueue(chunk);
      buffer += decoder.decode(chunk, { stream: true });
      const lines = buffer.split("\n");
      buffer = lines.pop() ?? "";
      for (const line of lines) observeLine(line);
    },
    flush(): void {
      buffer += decoder.decode();
      if (buffer) observeLine(buffer);
    },
  });
};