// What model calls cost, and a budget that stops spending at a set limit.
// Harnesses that know their prices put `costUsd` on each result's usage;
// the workspace charges every run and every grading to the budget.

import type { Harness } from "./harness.ts";
import type { Usage } from "./types.ts";

/** USD per million tokens. `cachedInput` applies to input served from the prompt cache. */
export type Price = { input: number; cachedInput?: number; output: number };

const p = (input: number, cachedInput: number, output: number): Price => ({ input, cachedInput, output });

/**
 * Tinker sampling prices (prefill, cached prefill, sample), from
 * https://tinker-docs.thinkingmachines.ai/tinker/models/ as of 2026-10-07.
 * The OpenAI-compatible endpoint samples at these rates. Tinker caches
 * prompt prefixes automatically; responses report the cached tokens.
 * Inkling's listed prices carry a limited-time 50% discount; the full
 * prices are used here, so the estimate errs high.
 */
export const TINKER_PRICES: Record<string, Price> = {
  "thinkingmachines/Inkling": p(3.74, 0.374, 9.36),
  "thinkingmachines/Inkling:peft:262144": p(7.48, 0.748, 18.72),
  "thinkingmachines/Inkling-Small": p(1.16, 0.116, 2.88),
  "thinkingmachines/Inkling-Small:peft:262144": p(2.32, 0.232, 5.78),
  "Qwen/Qwen3.8-27B": p(1.86, 0.372, 5.595),
  "Qwen/Qwen3.8-27B:peft:262144": p(2.48, 0.496, 7.46),
  "Qwen/Qwen3.6-35B-A3B": p(0.54, 0.108, 1.335),
  "Qwen/Qwen3.6-27B": p(1.86, 0.372, 5.595),
  "Qwen/Qwen3.5-397B-A17B": p(3.0, 0.6, 7.5),
  "Qwen/Qwen3.5-35B-A3B-Base": p(0.54, 0.108, 1.335),
  "Qwen/Qwen3.5-9B": p(0.66, 0.132, 1.995),
  "Qwen/Qwen3.5-9B-Base": p(0.66, 0.132, 1.995),
  "Qwen/Qwen3.5-4B": p(0.33, 0.066, 1.005),
  "Qwen/Qwen3-8B": p(0.195, 0.039, 0.6),
  "openai/gpt-oss-120b": p(0.33, 0.066, 0.84),
  "openai/gpt-oss-20b": p(0.18, 0.036, 0.45),
  "deepseek-ai/DeepSeek-V3.1": p(1.695, 0.339, 4.215),
  "moonshotai/Kimi-K2.6": p(2.205, 0.441, 5.49),
};

/** Tinker training prices (USD per million tokens), same source and date as TINKER_PRICES. */
export const TINKER_TRAIN_PRICES: Record<string, number> = {
  "thinkingmachines/Inkling": 11.22,
  "thinkingmachines/Inkling-Small": 3.46,
  "Qwen/Qwen3.8-27B": 4.103,
  "Qwen/Qwen3.6-35B-A3B": 1.177,
  "Qwen/Qwen3.6-27B": 4.103,
  "Qwen/Qwen3.5-397B-A17B": 6.6,
  "Qwen/Qwen3.5-35B-A3B-Base": 1.177,
  "Qwen/Qwen3.5-9B": 1.463,
  "Qwen/Qwen3.5-9B-Base": 1.463,
  "Qwen/Qwen3.5-4B": 0.737,
  "Qwen/Qwen3-8B": 0.44,
  "openai/gpt-oss-120b": 0.737,
  "openai/gpt-oss-20b": 0.396,
  "deepseek-ai/DeepSeek-V3.1": 3.718,
  "moonshotai/Kimi-K2.6": 4.84,
};

/** OpenRouter's free endpoints (`…:free`, e.g. thinkingmachines/inkling:free) cost nothing. */
export const FREE: Price = { input: 0, cachedInput: 0, output: 0 };
export function isFreeModel(model: string): boolean {
  return model.endsWith(":free");
}

const CHECKPOINT_BASES = new Map<string, string>();

/**
 * Record the base model a fine-tuned `tinker://` checkpoint was trained from,
 * so it is priced (and trained further) as that base. A machine's model.json
 * registers its checkpoint when the machine loads (see modelFrom).
 */
export function priceCheckpointAs(checkpoint: string, baseModel: string): void {
  if (!checkpoint.startsWith("tinker://")) return;
  if (!(baseModel in TINKER_PRICES)) throw new Error(`Unknown base model "${baseModel}" for ${checkpoint}: it has no Tinker price.`);
  CHECKPOINT_BASES.set(checkpoint, baseModel);
}

/** The base model of a registered checkpoint, the model itself for a base model, or undefined. */
export function baseModelOf(model: string): string | undefined {
  return model.startsWith("tinker://") ? CHECKPOINT_BASES.get(model) : model;
}

/** Price lookup for a Tinker model id. Fine-tuned checkpoints (`tinker://…`) are priced as their registered base, else `checkpointBase`. */
export function tinkerPrice(model: string, checkpointBase?: string): Price | undefined {
  if (model.startsWith("tinker://")) {
    const base = CHECKPOINT_BASES.get(model) ?? checkpointBase;
    return base ? TINKER_PRICES[base] : undefined;
  }
  return TINKER_PRICES[model];
}

export function costOf(price: Price, usage: Usage): number {
  const cached = Math.min(usage.cachedInputTokens ?? 0, usage.inputTokens ?? 0);
  const fresh = (usage.inputTokens ?? 0) - cached;
  return (fresh * price.input + cached * (price.cachedInput ?? price.input) + (usage.outputTokens ?? 0) * price.output) / 1_000_000;
}

/** Sum usages; the cost is known only if every part's cost is known. */
export function addUsage(...parts: Array<Usage | undefined>): Usage {
  const total: Usage = {};
  let costKnown = true;
  for (const u of parts) {
    if (!u) continue;
    if (u.inputTokens !== undefined) total.inputTokens = (total.inputTokens ?? 0) + u.inputTokens;
    if (u.outputTokens !== undefined) total.outputTokens = (total.outputTokens ?? 0) + u.outputTokens;
    if (u.cachedInputTokens !== undefined) total.cachedInputTokens = (total.cachedInputTokens ?? 0) + u.cachedInputTokens;
    if (u.costUsd !== undefined) total.costUsd = (total.costUsd ?? 0) + u.costUsd;
    else if (u.inputTokens || u.outputTokens) costKnown = false;
  }
  if (!costKnown) delete total.costUsd;
  return total;
}

/** A harness whose results carry their cost, for harnesses that report tokens but not prices. */
export function withPricing<H extends Harness>(harness: H, price: (model: string) => Price | undefined): H {
  return {
    ...harness,
    name: harness.name,
    async run(request) {
      const result = await harness.run(request);
      const p = result.usage && result.usage.costUsd === undefined ? price(request.model) : undefined;
      return p ? { ...result, usage: { ...result.usage, costUsd: costOf(p, result.usage!) } } : result;
    },
    ...(harness.close ? { close: () => harness.close!() } : {}),
  };
}

export class BudgetExceeded extends Error {
  constructor(message: string) {
    super(message);
    this.name = "BudgetExceeded";
  }
}

/**
 * A spending limit for one CLI invocation or script. `check()` before a model
 * call refuses once the limit is reached; `charge()` after it records the
 * spend. Calls already in flight finish, so the overrun is bounded by the
 * cost of the runs running at once.
 */
export class Budget {
  readonly limitUsd: number;
  spentUsd = 0;
  tokens = { input: 0, cachedInput: 0, output: 0 };

  constructor(limitUsd: number) {
    if (!(limitUsd >= 0)) throw new Error("A budget needs a non-negative limit in USD.");
    this.limitUsd = limitUsd;
  }

  check(what = "this call"): void {
    if (this.spentUsd >= this.limitUsd) {
      throw new BudgetExceeded(`Budget of $${this.limitUsd.toFixed(2)} reached ($${this.spentUsd.toFixed(4)} spent); refusing ${what}. Raise it with --budget.`);
    }
  }

  /** Record usage. Usage with tokens but no known cost is refused: the budget could not bound it. */
  charge(usage: Usage | undefined, what: string): void {
    if (!usage) return;
    this.tokens.input += usage.inputTokens ?? 0;
    this.tokens.output += usage.outputTokens ?? 0;
    this.tokens.cachedInput += usage.cachedInputTokens ?? 0;
    if (usage.costUsd !== undefined) this.spentUsd += usage.costUsd;
    else if (usage.inputTokens || usage.outputTokens) {
      throw new Error(`${what} used tokens at an unknown price, so the budget cannot bound it. Use a priced model or run without a budget (--budget none).`);
    }
  }

  summary(): string {
    return `spent $${this.spentUsd.toFixed(4)} of $${this.limitUsd.toFixed(2)} (${this.tokens.input} input of which ${this.tokens.cachedInput} cached, ${this.tokens.output} output tokens)`;
  }
}