import { Tiktoken } from 'js-tiktoken/lite';
import cl100k_base from 'js-tiktoken/ranks/cl100k_base';
import o200k_base from 'js-tiktoken/ranks/o200k_base';
 
const ENCODERS = new Map<TiktokenEncoding, Tiktoken>();
 
function getEncoder(name: TiktokenEncoding): Tiktoken {
  let enc = ENCODERS.get(name);
  if (enc) return enc;
  const ranks = name === 'o200k_base' ? o200k_base : cl100k_base;
  enc = new Tiktoken(ranks);
  ENCODERS.set(name, enc);
  return enc;
}
 
export function estimatePromptTokens(messages: IRMessage[], model: string): number {
  const enc = getEncoder(pickEncoding(model));
  let tokens = 0;
  for (const m of messages) {
    tokens += 4;                       // 消息封装固定开销 (role + 分隔符)
    tokens += enc.encode(m.role).length;
    if (typeof m.content === 'string') {
      tokens += enc.encode(m.content).length;
    } else if (Array.isArray(m.content)) {
      // multimodal: text 段累加; image_url 段固定 85 token (粗估)
      for (const part of m.content) {
        if (part?.type === 'text') tokens += enc.encode(part.text).length;
        else tokens += 85;
      }
    }
  }
  tokens += 2;                          // assistant 回复引导 token (粗估)
  return tokens;
}

// src/db/schema.ts
export const prices = sqliteTable('prices', {
  id: integer('id').primaryKey({ autoIncrement: true }),
  model: text('model').notNull(),
  provider: text('provider').notNull(),
  /** input 单价: 微元 / 1M tokens. */
  inputPriceMicroPer1M: integer('input_price_micro_per_1m').notNull(),
  /** output 单价: 微元 / 1M tokens. */
  outputPriceMicroPer1M: integer('output_price_micro_per_1m').notNull(),
  /** 模型自身倍率 (千分位). 例: 某模型成本被运营定为基线 1.2x, 这里写 1200 */
  modelMultiplier: integer('model_multiplier').notNull().default(1000),
  effectiveFrom: integer('effective_from').notNull(),
  effectiveTo: integer('effective_to'),  // null = 仍有效
  createdAt: integer('created_at').notNull(),
});

cost_micro = tokens × inputPriceMicroPer1M / 1_000_000

export function getCurrentPrice(model: string, provider: string): ResolvedPrice {
  const cached = CACHE.get(`${model}::${provider}`);
  if (cached && cached.expiresAt > Date.now()) return cached.price;
 
  const rows = db
    .select()
    .from(prices)
    .where(
      and(
        eq(prices.model, model),
        eq(prices.provider, provider),
        lte(prices.effectiveFrom, Date.now()),
        or(isNull(prices.effectiveTo), gt(prices.effectiveTo, Date.now())),
      ),
    )
    .orderBy(desc(prices.effectiveFrom))
    .limit(1)
    .all();
 
  if (rows.length === 0) throw new PriceNotFoundError(model, provider);
  // ... 缓存写入 ...
}

final_cost = base_cost
           × (user_multiplier   / 1000)   // 用户折扣 / 加价
           × (channel_multiplier / 1000)  // 渠道加价 (Ch8 渠道池接入)
           × (model_multiplier   / 1000)  // 模型倍率 (运营定的)

export function resolveMultiplier(ctx: MultiplierContext): CombinedMultiplier {
  const userRows = db
    .select({ m: users.userMultiplier })
    .from(users)
    .where(eq(users.id, ctx.userId))
    .all();
  const userMul = userRows.length > 0 ? userRows[0]!.m : MULTIPLIER_SCALE;
 
  const price = getCurrentPrice(ctx.model, ctx.provider);
  const modelMul = price.modelMultiplier;
 
  // v0.5: 渠道倍率写死 1000. Ch8 渠道池接入后从 channels 表读
  const channelMul = MULTIPLIER_SCALE;
 
  const combined = userMul * channelMul * modelMul;
  return {
    user: userMul,
    channel: channelMul,
    model: modelMul,
    combinedScale1e9: combined,
    combinedFloat: combined / (MULTIPLIER_SCALE * MULTIPLIER_SCALE * MULTIPLIER_SCALE),
  };
}

preConsume:
  1. 用 tiktoken 算 estimated_prompt_tokens
  2. 用 request.max_tokens (或 DEFAULT_MAX_TOKENS=4096) 当 output 上限
  3. 按当前价格 + 倍率, 算出预扣金额 (向上取整, 防止四舍五入少扣 1 微元)
  4. 用乐观锁扣余额: UPDATE WHERE balance >= cost
     失败 = 余额不够, 返 402 不进上游
  5. 写一行 usage_records, status = reserved

(主路径调上游)

postConsume / refundReservation (二选一):
  - 调上游成功: 按真实 usage 重算 final_cost, 多退少补
                UPDATE balance += (preReserved - finalCost)
                usage_records SET status = finalized
  - 调上游失败 (网络错 / 业务错): 全额退回预扣
                UPDATE balance += preReserved
                usage_records SET status = refunded

export function preConsume(input: PreConsumeInput): PreConsumeOutput {
  // 1. 本地估算 prompt tokens
  const estimatedPromptTokens = estimatePromptTokens(input.messages, input.model);
 
  // 2. 取当前价格 + 三合一倍率
  const price = getCurrentPrice(input.model, input.provider);
  const mul = resolveMultiplier({
    userId: input.userId, model: input.model, provider: input.provider,
  });
 
  // 3. 预扣金额 = (prompt × inputPrice + max_output × outputPrice) × 倍率
  const baseCost =
    estimatedPromptTokens * price.inputPricePerToken +
    input.maxOutputTokens * price.outputPricePerToken;
  const preReservedCost = Math.ceil(baseCost * mul.combinedFloat);
 
  // 4. 乐观锁扣余额: UPDATE WHERE balance >= cost
  const updated = db
    .update(users)
    .set({ balanceMicro: sql`${users.balanceMicro} - ${preReservedCost}` })
    .where(sql`${users.id} = ${input.userId} AND ${users.balanceMicro} >= ${preReservedCost}`)
    .returning({ balance: users.balanceMicro })
    .all();
  if (updated.length === 0) {
    throw new InsufficientBalanceError(preReservedCost, /* ... */);
  }
 
  // 5. 写入 usage_records 占位 (status = reserved)
  const rows = db.insert(usageRecords).values({
    traceId: input.traceId,
    userId: input.userId, orgId: input.orgId, keyId: input.keyId,
    model: input.model, provider: input.provider,
    estimatedPromptTokens,
    preReservedCost,
    multiplierSnapshot: mul.combinedScale1e9,
    status: 'reserved',
    isStream: input.isStream,
    createdAt: Date.now(),
  }).returning({ id: usageRecords.id }).all();
 
  return { recordId: rows[0]!.id, preReservedCost, estimatedPromptTokens };
}

export function postConsume(input: PostConsumeInput): PostConsumeOutput {
  const price = getCurrentPrice(input.model, input.provider);
  const mul = resolveMultiplier({ userId: input.userId, ... });
 
  const promptCost = Math.ceil(input.realPromptTokens * price.inputPricePerToken * mul.combinedFloat);
  const completionCost = Math.ceil(input.realCompletionTokens * price.outputPricePerToken * mul.combinedFloat);
  const finalCost = promptCost + completionCost;
 
  // 多退少补: balanceDelta = preReserved - finalCost
  //   > 0: 退回, balance += delta
  //   < 0: 少扣的部分再扣
  const balanceDelta = row.preReservedCost - finalCost;
  if (balanceDelta !== 0) {
    db.update(users)
      .set({ balanceMicro: sql`${users.balanceMicro} + ${balanceDelta}` })
      .where(eq(users.id, input.userId)).run();
  }
 
  db.update(usageRecords).set({
    promptTokens: input.realPromptTokens,
    completionTokens: input.realCompletionTokens,
    promptCost, completionCost, finalCost,
    status: 'finalized', finalizedAt: Date.now(),
  }).where(eq(usageRecords.id, input.recordId)).run();
 
  return { recordId, promptCost, completionCost, finalCost, balanceDelta };
}

export function refundReservation(recordId: number, errorMessage?: string): void {
  // ... 取 record ...
  if (row.status !== 'reserved') return;  // 幂等
 
  db.update(users)
    .set({ balanceMicro: sql`${users.balanceMicro} + ${row.preReservedCost}` })
    .where(eq(users.id, row.userId)).run();
  db.update(usageRecords).set({
    status: 'refunded', finalCost: 0, finalizedAt: Date.now(),
    errorMessage: errorMessage ?? null,
  }).where(eq(usageRecords.id, recordId)).run();
}

// src/index.ts (主路径骨架, 完整入参见仓库)
app.post('/v1/chat/completions', requireGatewayKey, async (c) => {
  const auth = c.get('auth');
  const ir = /* ... zod 校验 ... */;
  const adaptor = router.resolve(ir.model);
 
  // ----- 1) preConsume: 估 token + 扣余额 + 写 reserved 记录 -----
  let reservation;
  try {
    reservation = preConsume({ /* traceId / auth / ir / adaptor */ });
  } catch (err) {
    if (err instanceof InsufficientBalanceError) return c.json({ error: { type: 'insufficient_quota' } }, 402);
    if (err instanceof PriceNotFoundError) return c.json({ error: { type: 'price_not_configured' } }, 400);
    throw err;
  }
 
  // ----- 2) 调上游: 网络错 / 业务错都触发 refund -----
  let upstreamResp: Response;
  try {
    upstreamResp = await fetch(endpoint, { method: 'POST', headers, body });
  } catch (err) {
    refundReservation(reservation.recordId, `network_error: ${err.message}`);
    return c.json({ error: { message: 'upstream network error' } }, 502);
  }
  const rawBody = await upstreamResp.text();
  if (!upstreamResp.ok) {
    refundReservation(reservation.recordId, `upstream_${upstreamResp.status}`);
    return new Response(rawBody, { status: upstreamResp.status });
  }
 
  // ----- 3) postConsume: 用上游 usage (或本地 fallback) 实结 -----
  const irResponse = await adaptor.parseResponse(upstreamResp, rawBody);
  try {
    postConsume({
      recordId: reservation.recordId,
      realPromptTokens: irResponse.usage?.prompt_tokens ?? reservation.estimatedPromptTokens,
      realCompletionTokens: irResponse.usage?.completion_tokens
        ?? estimateCompletionTokens(irResponse.choices?.[0]?.message?.content ?? '', ir.model),
      /* userId / model / provider */
    });
  } catch (err) {
    markFailed(reservation.recordId, err.message);
  }
  return c.json(irResponse, 200);
});

export const usageRecords = sqliteTable('usage_records', {
  id: integer('id').primaryKey({ autoIncrement: true }),
  traceId: text('trace_id').notNull(),     // unique index, 32 字符 hex
  userId: integer('user_id').notNull(),
  orgId: integer('org_id').notNull(),
  keyId: integer('key_id').notNull(),
  model: text('model').notNull(),
  provider: text('provider').notNull(),
  promptTokens: integer('prompt_tokens').notNull().default(0),
  completionTokens: integer('completion_tokens').notNull().default(0),
  estimatedPromptTokens: integer('estimated_prompt_tokens').notNull().default(0),
  promptCost: integer('prompt_cost').notNull().default(0),
  completionCost: integer('completion_cost').notNull().default(0),
  finalCost: integer('final_cost').notNull().default(0),
  preReservedCost: integer('pre_reserved_cost').notNull().default(0),
  multiplierSnapshot: integer('multiplier_snapshot').notNull().default(1_000_000_000),
  status: text('status').notNull().default('reserved'),
  isStream: integer('is_stream', { mode: 'boolean' }).notNull().default(false),
  errorMessage: text('error_message'),
  createdAt: integer('created_at').notNull(),
  finalizedAt: integer('finalized_at'),
});

export class StreamingTokenCounter {
  private upstreamPromptTokens: number | null = null;
  private upstreamCompletionTokens: number | null = null;
  private localCompletionText = '';
  private localCompletionTokens = 0;
  private aborted = false;
 
  constructor(model: string, fallbackPromptTokens: number) {
    this.model = model;
    this.fallbackPromptTokens = fallbackPromptTokens;
  }
 
  /** 收到一段 completion 文本 delta. 累加本地估算 */
  ingestDelta(deltaText: string): void {
    if (!deltaText) return;
    this.localCompletionText += deltaText;
    this.localCompletionTokens += estimateCompletionTokens(deltaText, this.model);
  }
 
  /** 上游送了 usage 事件 (OpenAI stream_options.include_usage / Anthropic message_delta.usage) */
  ingestUsage(usage: { prompt_tokens?: number; completion_tokens?: number }): void {
    if (typeof usage.prompt_tokens === 'number') this.upstreamPromptTokens = usage.prompt_tokens;
    if (typeof usage.completion_tokens === 'number') this.upstreamCompletionTokens = usage.completion_tokens;
  }
 
  /** 客户端中途取消时调一次 */
  markAborted(): void {
    this.aborted = true;
  }
 
  /** 流结束 (含中断) 时返回最终 token 数, 喂给 postConsume */
  finalize(): StreamingFinalize {
    const promptTokens = this.upstreamPromptTokens ?? this.fallbackPromptTokens;
    const completionTokens = this.upstreamCompletionTokens ?? this.localCompletionTokens;
    return { promptTokens, completionTokens, abortedByClient: this.aborted };
  }
}

SELECT
  user_id,
  COUNT(*)                                 AS request_count,
  SUM(prompt_tokens)                       AS prompt_tokens,
  SUM(completion_tokens)                   AS completion_tokens,
  SUM(final_cost) / 1000000.0              AS final_cost_cny
FROM usage_records
WHERE status = 'finalized'
  AND created_at >= strftime('%s','now','start of day') * 1000
GROUP BY user_id
ORDER BY final_cost_cny DESC;

SELECT
  model, provider,
  COUNT(*)                                 AS request_count,
  SUM(prompt_tokens)                       AS prompt_tokens,
  SUM(completion_tokens)                   AS completion_tokens,
  SUM(prompt_cost) / 1000000.0             AS prompt_cost_cny,
  SUM(completion_cost) / 1000000.0         AS completion_cost_cny,
  SUM(final_cost) / 1000000.0              AS total_cny
FROM usage_records
WHERE status = 'finalized'
  AND created_at >= strftime('%s','now','start of month') * 1000
GROUP BY model, provider
ORDER BY total_cny DESC;

SELECT
  strftime('%Y-%m-%d', datetime(created_at / 1000, 'unixepoch')) AS day,
  COUNT(*)                                 AS request_count,
  SUM(final_cost) / 1000000.0              AS daily_cost_cny
FROM usage_records
WHERE status = 'finalized'
  AND created_at >= (strftime('%s','now') - 30 * 86400) * 1000
GROUP BY day
ORDER BY day DESC;

cd examples/05-who-spent-my-money
cp .env.example .env
# 改 ADMIN_TOKEN; 上游 Key 暂时留 sk-replace-me (本章主要看计费链路, 上游失败属于预期)
npm install
npm run dev

INFO db_migrations_applied applied=["0001_init.sql","0002_billing.sql"]
INFO default_prices_seeded inserted=10
INFO Gateway v0.5 listening on http://localhost:3000

ADMIN_TOKEN=test-admin-token-12345
BASE=http://localhost:3000
 
curl -s -X POST $BASE/admin/orgs -H "Authorization: Bearer $ADMIN_TOKEN" \
  -H "Content-Type: application/json" -d '{"name":"Acme Inc"}'
# {"id":1,"name":"Acme Inc","disabledAt":null,"createdAt":...}
 
curl -s -X POST $BASE/admin/users -H "Authorization: Bearer $ADMIN_TOKEN" \
  -H "Content-Type: application/json" -d '{"orgId":1,"name":"alice","balanceCny":100}'
# {"id":1,"orgId":1,"name":"alice","balanceMicro":100000000,"userMultiplier":1000,...}
 
KEY=$(curl -s -X POST $BASE/admin/keys -H "Authorization: Bearer $ADMIN_TOKEN" \
  -H "Content-Type: application/json" -d '{"userId":1,"name":"smoke"}' \
  | python3 -c "import json,sys;print(json.load(sys.stdin)['plaintext'])")

curl -s -X POST $BASE/v1/chat/completions \
  -H "Authorization: Bearer $KEY" \
  -H "Content-Type: application/json" \
  -d '{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hello world, this is a test message for billing"}]}'
# {"error":{"message":"upstream network error"}}

curl -s "$BASE/admin/usage?userId=1" -H "Authorization: Bearer $ADMIN_TOKEN"

{
  "data": [{
    "id": 1,
    "traceId": "dfa123480dcd949104eb0e6c18f74672",
    "userId": 1, "orgId": 1, "keyId": 1,
    "model": "gpt-4o-mini", "provider": "openai",
    "promptTokens": 0, "completionTokens": 0,
    "estimatedPromptTokens": 17,
    "promptCost": 0, "completionCost": 0, "finalCost": 0,
    "preReservedCost": 17713,
    "multiplierSnapshot": 1000000000,
    "status": "refunded",
    "isStream": false,
    "errorMessage": "network_error: fetch failed",
    "createdAt": ..., "finalizedAt": ...
  }]
}

SELECT id, name, balance_micro FROM users;
-- 1 | alice | 100000000  (一亿微元 = 100 元, 没动)

curl -s -X POST $BASE/admin/users/1/balance -H "Authorization: Bearer $ADMIN_TOKEN" \
  -H "Content-Type: application/json" -d '{"setCny":0.001}'
 
curl -i -s -X POST $BASE/v1/chat/completions \
  -H "Authorization: Bearer $KEY" \
  -H "Content-Type: application/json" \
  -d '{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}'

HTTP/1.1 402 Payment Required
{"error":{"type":"insufficient_quota","message":"balance is not enough to cover the reservation","required_micro_cny":17704,"available_micro_cny":1000}}

{
  "promptTokens": 17,
  "completionTokens": 23,
  "estimatedPromptTokens": 17,
  "promptCost": 18,
  "completionCost": 99,
  "finalCost": 117,
  "preReservedCost": 17713,
  "status": "finalized",
  "finalizedAt": 1715750999999
}

SELECT key_id, COUNT(*) AS req, MIN(created_at) AS first_seen, MAX(created_at) AS last_seen
FROM usage_records
WHERE created_at >= strftime('%s','now','-1 day') * 1000
GROUP BY key_id
HAVING req > 10000
ORDER BY req DESC;

src/
  index.ts                  # Hono 入口, 主路径接入两阶段计费
  db/
    schema.ts               # (new) users 加 balance + multiplier; 新增 prices / usage_records
    migrate.ts              # 与 v0.4 一致
    client.ts               # 与 v0.4 一致
  billing/                  # (new) 本章新增
    prices.ts               # 价格表 + TTL 缓存 + 默认价灌库
    tokenizer.ts            # js-tiktoken 本地估算
    calculator.ts           # preConsume / postConsume / refund 三段式
    streaming-counter.ts    # 流式 token 计数器 (Ch7 接入 SSE 主循环)
    record.ts               # UsageRecord 查询辅助
  multiplier/
    registry.ts             # (new) user × channel × model 三合一倍率
  auth/                     # 与 v0.4 一致
  admin/
    routes.ts               # (new) 扩展: 调余额 / 改倍率 / 列价格 / 改价 / 查账单
  cli/
    issue-key.ts            # (new) 改动: 创建 user 时给初始余额
  adaptors/                 # 与 v0.4 一致
  streaming/                # 与 v0.4 一致
  types/                    # 与 v0.4 一致
  router.ts                 # 与 v0.4 一致
drizzle/
  0001_init.sql             # 与 v0.4 一致
  0002_billing.sql          # (new) 本章新增