/**
 * dsh-budget — spend ENFORCEMENT for DeepSeek Harness.
 *
 * The ecosystem is full of plugins that measure token spend after the fact.
 * This one sits on the `llm/stream` waterfall, which is a *gate*: the decision
 * happens BEFORE `next()` is called, so an over-budget request is never billed.
 * Usage chunks emitted by the adapter are metered on the way back out.
 * @module dsh-budget
 */
import type { Context } from '@deepseek-ai/cordis'
// Importing the llm package applies its module augmentation, which is what
// declares the `llm/stream` waterfall on cordis' Events map.
import type { GenerateOptions, StreamChunk } from '@deepseek-ai/dsh-llm'
// @ts-expect-error -- zero-dependency JS core, typed via ledger.d.ts
import { SpendStore, priceCall, periodKey, decide, resolveScope, limitFor, MICRO } from '../src/ledger.mjs'

export interface ModelRates {
  /** USD per MILLION tokens, as providers publish them. */
  input: number
  output: number
  cacheRead?: number
  cacheWrite?: number
  reasoning?: number
}

export interface BudgetConfig {
  /** Hard spend cap in USD for the period. 0 or absent disables enforcement. */
  limitUsd?: number
  /** Budget window. */
  period?: 'daily' | 'monthly' | 'total'
  /** Fraction of the limit at which to start downgrading. */
  softFraction?: number
  /** Cheaper model used once past `softFraction`. */
  downgradeModel?: string
  /** Per-model pricing, keyed by model id. */
  rates?: Record<string, ModelRates>
  /** Where the ledger persists. */
  storePath?: string
  /** Report but never block or downgrade — for rolling out safely. */
  dryRun?: boolean
  /**
   * Attribution scope. `deployment` (default) is one budget for the whole
   * install. `session` uses `GenerateOptions.sessionId`, which dsh provides.
   * `user` needs a mapping dsh cannot supply — see `userEnv`/`sessionUsers`.
   */
  scopeBy?: 'deployment' | 'session' | 'user'
  /** Env var holding the current user — for one-instance-per-user deployments. */
  userEnv?: string
  /** Operator-supplied `{ sessionId: user }` map. Wins over `userEnv`. */
  sessionUsers?: Record<string, string>
  /** Per-scope limit overrides in USD. `null` means unlimited. */
  perScopeLimitUsd?: Record<string, number | null>
}

export const name = 'dsh-cost-gate'

export function apply(ctx: Context, config: BudgetConfig = {}): void {
  const limitMicros = Math.round((config.limitUsd ?? 0) * MICRO)
  const period = config.period ?? 'monthly'
  const store = new SpendStore(config.storePath ?? null)
  const rates = config.rates ?? {}

  const perScopeLimitMicros = Object.fromEntries(
    Object.entries(config.perScopeLimitUsd ?? {}).map(([k, v]) => [k, v === null ? null : Math.round(v * MICRO)]),
  )

  ctx.on('llm/stream', function (options: GenerateOptions, next: () => AsyncIterable<StreamChunk>) {
    const key = periodKey(period, Date.now())
    // Internal calls (compaction, session titles) are system overhead and get
    // their own bucket, so they can never exhaust a person's budget.
    const { scope } = resolveScope(options, config)
    const scopeLimit = limitFor(scope, { limitMicros, perScopeLimitMicros })
    const spentMicros = store.get(scope, key)
    const verdict = decide({
      spentMicros, limitMicros: scopeLimit,
      softFraction: config.softFraction ?? 0.8,
      downgradeModel: config.downgradeModel,
      requestedModel: options.model,
    })

    if (verdict.action === 'block' && !config.dryRun) {
      ctx.logger?.warn?.(`[dsh-cost-gate] BLOCKED ${scope}: ${verdict.reason}`)
      throw new Error(
        `dsh-cost-gate: ${scope} ${verdict.reason}. Raise limitUsd, wait for the ${period} reset, or set dryRun.`,
      )
    }
    if (verdict.action === 'downgrade' && !config.dryRun) {
      ctx.logger?.info?.(`[dsh-cost-gate] ${scope}: ${verdict.reason}`)
      ;(options as { model: string }).model = verdict.model as string
    }
    if (config.dryRun && verdict.action !== 'allow') {
      ctx.logger?.info?.(`[dsh-cost-gate] dry-run would ${verdict.action} ${scope}: ${verdict.reason}`)
    }

    // Meter on the way out: usage chunks arrive before the terminal finish.
    const modelForBilling = options.model
    const upstream = next()
    return (async function* () {
      for await (const chunk of upstream) {
        const c = chunk as unknown as { type?: string; usage?: Record<string, number> }
        if (c.type === 'usage' && c.usage) {
          const micros = priceCall(c.usage, rates[modelForBilling])
          if (micros > 0) store.add(scope, key, micros)
        }
        yield chunk
      }
    })()
  })
}

export default apply