import type { Message, AudioSource } from '../models/types.js';
import { getMessageText } from '../models/media.js';

/**
 * Approximate overhead tokens per message (role marker, separators, etc.).
 */
const MESSAGE_OVERHEAD_TOKENS = 4;

/**
 * Conservative default context window when a model is not in the lookup table.
 */
const DEFAULT_CONTEXT_WINDOW = 128_000;

/**
 * Hard-coded context window sizes (in tokens) for known models.
 */
export const CONTEXT_WINDOWS: Record<string, number> = {
  'claude-sonnet-4-20250514': 200_000,
  'claude-3-5-haiku-20241022': 200_000,
  'claude-3-5-sonnet-20241022': 200_000,
  'claude-3-opus-20240229': 200_000,
  'claude-opus-4-20250514': 200_000,
  'gpt-4o': 128_000,
  'gpt-4o-mini': 128_000,
  'gpt-4-turbo': 128_000,
  'gpt-3.5-turbo': 16_385,
} as const;

/**
 * Cheap character-based token estimation.
 *
 * Uses `Math.ceil(text.length / 4)` as a reasonable approximation for
 * English text (roughly 4 characters per token on average).
 */
export function estimateTokens(text: string): number {
  return Math.ceil(text.length / 4);
}

/**
 * Estimate token count for an audio content part.
 *
 * Heuristic:
 *  1. Decode base64 length to bytes: `base64Length * 0.75`
 *  2. Assume ~16 kbps bitrate (typical voice OGG/Opus): `bytes / 2000` → seconds
 *  3. Estimate ~32 tokens per second of audio (Gemini-style rate)
 *
 * Returns at least 1 token for any non-empty audio data.
 */
export function estimateAudioTokens(audioSource: AudioSource): number {
  const base64Length = audioSource.data.length;
  if (base64Length === 0) {
    return 0;
  }
  const durationSeconds = (base64Length * 0.75) / 2000;
  return Math.max(1, Math.ceil(durationSeconds * 32));
}

/**
 * Estimate the total token count for an array of messages.
 *
 * For each message the estimate includes the content tokens plus a fixed
 * overhead of ~4 tokens to account for the role marker and separators.
 */
export function estimateMessageTokens(messages: Message[]): number {
  return messages.reduce((sum, msg) => {
    let tokens = estimateTokens(getMessageText(msg)) + MESSAGE_OVERHEAD_TOKENS;

    // Add audio token estimates for multimodal messages
    if (Array.isArray(msg.content)) {
      for (const part of msg.content) {
        if (part.type === 'audio') {
          tokens += estimateAudioTokens(part.source);
        }
      }
    }

    return sum + tokens;
  }, 0);
}

/**
 * Return the context window size (in tokens) for a given model.
 *
 * @param model    - Model identifier to look up.
 * @param override - If provided, this value is returned directly.
 * @returns The context window size in tokens.
 */
export function getContextWindow(model: string, override?: number): number {
  if (override !== undefined) {
    return override;
  }
  return CONTEXT_WINDOWS[model] ?? DEFAULT_CONTEXT_WINDOW;
}

/**
 * Options for {@link shouldCompact}.
 */
export interface ShouldCompactOpts {
  messages: Message[];
  model: string;
  contextWindow?: number;
  /** Percentage of the context window that triggers compaction (default 80). */
  thresholdPct?: number;
}

/**
 * Determine whether the conversation should be compacted.
 *
 * Returns `true` when the estimated token count of `messages` exceeds
 * `thresholdPct` percent of the effective context window.
 */
export function shouldCompact(opts: ShouldCompactOpts): boolean {
  const { messages, model, contextWindow, thresholdPct = 80 } = opts;
  const window = getContextWindow(model, contextWindow);
  const threshold = (thresholdPct / 100) * window;
  const estimated = estimateMessageTokens(messages);
  return estimated > threshold;
}
