import type { TokenUsage } from '#cli/core/model'

type MutableTokenUsage = { -readonly [K in keyof TokenUsage]: TokenUsage[K] }

type UsageState = {
  chromeVersion: number
  usage: MutableTokenUsage
}

type UsagePersistence = Readonly<{
  applyUsageFromSession: (usage: unknown) => void
}>

export function createSessionUsagePersistence(
  opts: Readonly<{
    state: UsageState
  }>,
): UsagePersistence {
  const { state } = opts

  function applyUsageFromSession(usage: unknown): void {
    if (!usage || typeof usage !== 'object')
      return

    const u = usage as any
    const prompt = Number(u.promptTokens ?? u.prompt_tokens)
    const completion = Number(u.completionTokens ?? u.completion_tokens)
    const total = Number(u.totalTokens ?? u.total_tokens)
    const normalizedPrompt = Number.isFinite(prompt)
      ? Math.max(0, prompt)
      : Number.NaN
    const normalizedCompletion = Number.isFinite(completion)
      ? Math.max(0, completion)
      : Number.NaN

    // Preserve the last non-zero request usage so failed/partial updates do not
    // clear previously displayed values.
    let nextPrompt = state.usage.prompt_tokens
    let nextCompletion = state.usage.completion_tokens
    let nextTotal = state.usage.total_tokens
    if (
      Number.isFinite(normalizedPrompt)
      && normalizedPrompt > 0
      && nextPrompt !== normalizedPrompt
    ) {
      nextPrompt = normalizedPrompt
    }
    if (
      Number.isFinite(normalizedCompletion)
      && normalizedCompletion > 0
      && nextCompletion !== normalizedCompletion
    ) {
      nextCompletion = normalizedCompletion
    }

    // Keep total monotonically non-decreasing. Prefer SDK cumulative total,
    // with prompt+completion as fallback when total is unavailable.
    const totalCandidate = Number.isFinite(total)
      ? Math.max(0, total)
      : Number.isFinite(normalizedPrompt)
        && Number.isFinite(normalizedCompletion)
        ? normalizedPrompt + normalizedCompletion
        : Number.NaN
    const currentTotal = Math.max(0, Number(state.usage.total_tokens) || 0)
    if (
      Number.isFinite(totalCandidate)
      && totalCandidate >= currentTotal
      && nextTotal !== totalCandidate
    ) {
      nextTotal = totalCandidate
    }
    const changed
      = nextPrompt !== state.usage.prompt_tokens
        || nextCompletion !== state.usage.completion_tokens
        || nextTotal !== state.usage.total_tokens
    if (changed) {
      state.usage = {
        ...state.usage,
        prompt_tokens: nextPrompt,
        completion_tokens: nextCompletion,
        total_tokens: nextTotal,
      }
    }
    if (changed)
      state.chromeVersion += 1
  }

  return {
    applyUsageFromSession,
  }
}
