import type { PolyModel } from "@/lib/client/types";

/**
 * Client-side token & cost estimation. Heuristic (≈ 4 chars/token for English prose, ~3.2 for code)
 * — precise enough for "≈ $0.08" hints; the server records exact provider usage afterwards.
 */
export function estimateTextTokens(text: string): number {
  if (!text) return 0;
  const codeish = /[{}();=<>[\]]/.test(text) && (text.match(/\n/g)?.length ?? 0) > 2;
  const perToken = codeish ? 3.2 : 4.0;
  // CJK scripts are ~1 token per character.
  const cjk = (text.match(/[぀-ヿ㐀-䶿一-鿿가-힯]/g) ?? []).length;
  return Math.ceil((text.length - cjk) / perToken + cjk);
}

export interface AttachmentLike {
  kind: string;
  sizeBytes: number;
  width?: number | null;
  height?: number | null;
}

/** Rough per-attachment token cost (images follow the common 512px-tile heuristic; text ≈ bytes / 4; PDFs ≈ 1.5K / page ≈ bytes/45). */
export function estimateAttachmentTokens(a: AttachmentLike): number {
  if (a.kind === "image") {
    const w = a.width ?? 1024;
    const h = a.height ?? 1024;
    const tiles = Math.ceil(Math.min(w, 2048) / 512) * Math.ceil(Math.min(h, 2048) / 512);
    return 85 + tiles * 170;
  }
  if (a.kind === "pdf") return Math.max(300, Math.ceil(a.sizeBytes / 45));
  return Math.ceil(a.sizeBytes / 4);
}

export interface ContextEstimate {
  /** Tokens in the history that will be replayed. */
  history: number;
  /** Tokens in the draft + attachments + system prompt. */
  draft: number;
  total: number;
  contextTokens: number | null;
  /** 0..1 or null when the model has no known limit. */
  ratio: number | null;
  level: "ok" | "warn" | "critical" | "over";
}

export function estimateContext(opts: { historyText: string[]; historyAttachments?: AttachmentLike[]; draft: string; attachments?: AttachmentLike[]; systemPrompt?: string | null; model?: PolyModel | null; expectedOutput?: number }): ContextEstimate {
  const history = opts.historyText.reduce((n, t) => n + estimateTextTokens(t), 0) + (opts.historyAttachments ?? []).reduce((n, a) => n + estimateAttachmentTokens(a), 0) + opts.historyText.length * 4;
  const draft = estimateTextTokens(opts.draft) + (opts.attachments ?? []).reduce((n, a) => n + estimateAttachmentTokens(a), 0) + estimateTextTokens(opts.systemPrompt ?? "");
  const total = history + draft;
  const contextTokens = opts.model?.limits?.contextTokens ?? null;
  const ratio = contextTokens ? (total + (opts.expectedOutput ?? 0)) / contextTokens : null;
  const level = ratio === null ? "ok" : ratio >= 1 ? "over" : ratio >= 0.95 ? "critical" : ratio >= 0.8 ? "warn" : "ok";
  return { history, draft, total, contextTokens, ratio, level };
}

export interface CostEstimate {
  inputTokens: number;
  outputTokens: number;
  /** null when the model has no public pricing. */
  usd: number | null;
  inputUsd: number | null;
  outputUsd: number | null;
}

/** Estimate a request cost for a model: input tokens known, output tokens assumed (default 600, or maxTokens when smaller). */
export function estimateCost(model: PolyModel | null | undefined, inputTokens: number, outputTokens = 600): CostEstimate {
  const p = model?.pricing;
  if (!p || (p.inputPerMillion === undefined && p.outputPerMillion === undefined)) return { inputTokens, outputTokens, usd: null, inputUsd: null, outputUsd: null };
  let inRate = p.inputPerMillion ?? 0;
  let outRate = p.outputPerMillion ?? 0;
  if (p.longContext && inputTokens > p.longContext.thresholdTokens) {
    inRate = p.longContext.inputPerMillion ?? inRate;
    outRate = p.longContext.outputPerMillion ?? outRate;
  }
  const inputUsd = (inputTokens / 1_000_000) * inRate;
  const outputUsd = (outputTokens / 1_000_000) * outRate;
  return { inputTokens, outputTokens, usd: inputUsd + outputUsd, inputUsd, outputUsd };
}

/** Compact "~22K tokens · ≈ $0.08" label. */
export function formatEstimate(tokens: number, usd: number | null): string {
  const t = tokens >= 1_000_000 ? `${(tokens / 1_000_000).toFixed(1)}M` : tokens >= 1000 ? `${Math.round(tokens / 1000)}K` : String(tokens);
  if (usd === null) return `~${t} tokens`;
  const u = usd < 0.01 ? `$${usd.toFixed(4)}` : `$${usd.toFixed(2)}`;
  return `~${t} tokens · ≈ ${u}`;
}

/** Above this estimated cost the composer asks for confirmation before sending. */
export const COST_CONFIRM_THRESHOLD_USD = 1.0;