// TokenUsageDetector（veridrop detectors/token_usage.py）：短题、长题（多 80 个 " apple"）各问一次，再流式问一次短题、
// 调一次 count_tokens。五项：用量字段齐 20、长短输入差在范围内 25、输出不超 max_tokens 15、流式与非流式用量接近 25、
// count_tokens 与实际总输入接近 15。≥ 80 算过。
//
// count_tokens 这一项与原仓库（e87ed92）不同：原仓库拿它和 usage.input_tokens 单独比，用了提示缓存的请求（官方也一样）
// input_tokens 只剩缓存断点之后的几个 token，count_tokens 却数的是全部输入，比法本身就会误判。veridrop.app 线上版已经改成
// 和总输入（input_tokens + cache_read_input_tokens + cache_creation_input_tokens）比、容差总输入的 20%
// （它报告里的 actual_total_input_tokens / tolerance 字段），这里照线上版。

import { lookupModel } from "../config.js";
import type { DetectorResult } from "../models.js";
import { cpSlice, isInt, joinText } from "../py.js";
import { ActiveDetector, errText, type DetectClient } from "./base.js";

const SHORT_PROMPT = "Reply with exactly: ok";
const LONG_PROMPT = SHORT_PROMPT + "\n\nReference text:" + " apple".repeat(80);
const MAX_TOKENS = 16;
const DELTA_MIN = 45;
const DELTA_MAX = 140;
const NEW_TOKENIZER_DELTA_MIN = 90;
const NEW_TOKENIZER_DELTA_MAX = 230;
const OUTPUT_TOKEN_TOLERANCE = 4;
const INPUT_CLOSE_FRAC = 0.2;
const INPUT_CLOSE_FLOOR = 4;
const OUTPUT_CLOSE_FRAC = 0.5;
const OUTPUT_CLOSE_FLOOR = 3;

type Usage = Record<string, unknown>;

function intOrNull(v: unknown): number | null {
  return isInt(v) ? v : null;
}

function usageOf(resp: Record<string, unknown>): Usage {
  const u = resp["usage"];
  return u && typeof u === "object" && !Array.isArray(u) ? (u as Usage) : {};
}

export function deltaRange(model: string): [number, number] {
  const info = lookupModel(model);
  if (info !== null && info.newTokenizer) return [NEW_TOKENIZER_DELTA_MIN, NEW_TOKENIZER_DELTA_MAX];
  return [DELTA_MIN, DELTA_MAX];
}

function outputSane(usage: Usage): boolean {
  const out = intOrNull(usage["output_tokens"]);
  return out !== null && out >= 0 && out <= MAX_TOKENS + OUTPUT_TOKEN_TOLERANCE;
}

interface StreamUsage {
  input_tokens: number | null;
  output_tokens: number | null;
  chunk_count: number;
  error: string | null;
}

function streamUsageClose(ns: Usage, st: StreamUsage): boolean {
  const nsIn = intOrNull(ns["input_tokens"]);
  const nsOut = intOrNull(ns["output_tokens"]);
  if (nsIn === null || nsOut === null || st.input_tokens === null || st.output_tokens === null) return false;
  const inTol = Math.max(INPUT_CLOSE_FLOOR, Math.trunc(nsIn * INPUT_CLOSE_FRAC));
  const outTol = Math.max(OUTPUT_CLOSE_FLOOR, Math.trunc(nsOut * OUTPUT_CLOSE_FRAC));
  return Math.abs(nsIn - st.input_tokens) <= inTol && Math.abs(nsOut - st.output_tokens) <= outTol && st.output_tokens >= 0;
}

/** 实际总输入 = input + 缓存读 + 缓存写（缺的按 0）；input_tokens 本身没有回 null。 */
export function totalInput(ns: Usage): number | null {
  const input = intOrNull(ns["input_tokens"]);
  if (input === null) return null;
  return input + (intOrNull(ns["cache_read_input_tokens"]) ?? 0) + (intOrNull(ns["cache_creation_input_tokens"]) ?? 0);
}

export function countClose(ns: Usage, count: { input_tokens: number | null }): boolean {
  const actual = totalInput(ns);
  const counted = count.input_tokens;
  if (actual === null || counted === null) return false;
  const tol = Math.max(INPUT_CLOSE_FLOOR, Math.trunc(actual * INPUT_CLOSE_FRAC));
  return Math.abs(actual - counted) <= tol;
}

function evaluation(score: number, count: { error: string | null }): string {
  if (score >= 90) {
    return "Token 用量基本可信: usage 字段完整,长短 prompt 增量合理,stream 与 non-stream 统计一致,且 count_tokens 对得上。";
  }
  if (score >= 80) {
    if (count.error) {
      return "Token 用量基本可信: usage 自洽且流式/非流式一致,但该中转站没有可用的 count_tokens 端点,无法做最强输入 token 交叉验证。";
    }
    return "Token 用量基本可信: 核心 usage 检查通过,但有一个辅助交叉验证不完整。";
  }
  return "Token 用量存在风险: usage 字段缺失、长短 prompt 增量异常、输出 token 超出请求上限,或 stream 与 non-stream token 统计不一致。";
}

export class TokenUsageDetector extends ActiveDetector {
  readonly name = "token_usage";
  readonly displayName = "Token 用量";
  readonly weight = 10;

  async run(client: DetectClient, model: string): Promise<DetectorResult> {
    const shortBody = { model, max_tokens: MAX_TOKENS, temperature: 0, messages: [{ role: "user", content: SHORT_PROMPT }] };
    const longBody = { model, max_tokens: MAX_TOKENS, temperature: 0, messages: [{ role: "user", content: LONG_PROMPT }] };

    let shortResp: Record<string, unknown>;
    let longResp: Record<string, unknown>;
    let stream: StreamUsage;
    try {
      ({ resp: shortResp } = await client.messagesCreate(shortBody));
      ({ resp: longResp } = await client.messagesCreate(longBody));
      stream = await collectStreamUsage(client, shortBody);
    } catch (e) {
      return this.result("error", 0, null, null, errText(e));
    }
    const count = await countTokens(client, shortBody);

    const shortUsage = usageOf(shortResp);
    const longUsage = usageOf(longResp);
    const shortText = joinText(shortResp["content"]).trim();
    const longText = joinText(longResp["content"]).trim();

    const sub: Record<string, Record<string, unknown>> = {};
    let score = 0;

    // Python 的 bool(dict)：空 dict 算假。
    const usagePresent = Object.keys(shortUsage).length > 0 && Object.keys(longUsage).length > 0;
    sub["usage_present"] = { pass: usagePresent, short_usage: shortUsage, long_usage: longUsage };
    if (usagePresent) score += 20;

    const shortIn = intOrNull(shortUsage["input_tokens"]);
    const longIn = intOrNull(longUsage["input_tokens"]);
    const delta = shortIn !== null && longIn !== null ? longIn - shortIn : null;
    const [dMin, dMax] = deltaRange(model);
    const deltaOk = delta !== null && delta >= dMin && delta <= dMax;
    sub["input_token_delta"] = {
      short_input_tokens: shortIn,
      long_input_tokens: longIn,
      delta,
      expected_range: [dMin, dMax],
      pass: deltaOk,
    };
    if (deltaOk) score += 25;

    const outputOk = outputSane(shortUsage) && outputSane(longUsage);
    sub["output_tokens"] = {
      short_output_tokens: shortUsage["output_tokens"] ?? null,
      long_output_tokens: longUsage["output_tokens"] ?? null,
      max_tokens: MAX_TOKENS,
      short_text: cpSlice(shortText, 80),
      long_text: cpSlice(longText, 80),
      pass: outputOk,
    };
    if (outputOk) score += 15;

    const streamOk = streamUsageClose(shortUsage, stream);
    sub["stream_usage"] = {
      non_stream_usage: shortUsage,
      stream_input_tokens: stream.input_tokens,
      stream_output_tokens: stream.output_tokens,
      stream_chunk_count: stream.chunk_count,
      stream_error: stream.error,
      pass: streamOk,
    };
    if (streamOk) score += 25;

    const countOk = countClose(shortUsage, count);
    sub["count_tokens"] = {
      count_tokens: count.input_tokens,
      count_error: count.error,
      actual_input_tokens: shortIn,
      actual_total_input_tokens: totalInput(shortUsage),
      pass: countOk,
      note: countOk ? "count_tokens 与实际 Messages usage 接近" : "count_tokens 不可用或与实际 usage 偏差过大",
    };
    if (countOk) score += 15;

    if (!usagePresent) return this.skip("missing-usage");
    return this.result(score >= 80 ? "pass" : "fail", score, {
      sub_checks: sub,
      evaluation_zh: evaluation(score, count),
      risk_level: score >= 90 ? "low" : score >= 80 ? "medium" : "high",
    });
  }
}

/** 流式问一次：message_start 取 input_tokens，message_delta 取（最后一个）output_tokens。流中途出错不抛，记在 error 里。 */
async function collectStreamUsage(client: DetectClient, body: Record<string, unknown>): Promise<StreamUsage> {
  let input: number | null = null;
  let output: number | null = null;
  let chunks = 0;
  try {
    for await (const [ev] of client.messagesStream(body)) {
      chunks++;
      if (ev.event === "message_start") {
        const msg = (ev.data["message"] ?? {}) as Record<string, unknown>;
        input = intOrNull(((msg["usage"] ?? {}) as Record<string, unknown>)["input_tokens"]);
      } else if (ev.event === "message_delta") {
        const out = intOrNull(((ev.data["usage"] ?? {}) as Record<string, unknown>)["output_tokens"]);
        if (out !== null) output = out;
      }
    }
  } catch (e) {
    return { input_tokens: input, output_tokens: output, chunk_count: chunks, error: errText(e) };
  }
  return { input_tokens: input, output_tokens: output, chunk_count: chunks, error: null };
}

async function countTokens(
  client: DetectClient,
  body: Record<string, unknown>,
): Promise<{ input_tokens: number | null; error: string | null }> {
  try {
    const { resp } = await client.countTokens(body);
    return { input_tokens: intOrNull(resp["input_tokens"]), error: null };
  } catch (e) {
    return { input_tokens: null, error: errText(e) };
  }
}