/**
 * Model adapter registry. The app is not coupled to a specific model; any
 * Transformers.js-compatible ONNX decoder can be plugged in here.
 */

export interface ModelSpec {
  id: string;
  label: string;
  repo: string; // HF hub repo with ONNX weights
  dtype: 'q4f16' | 'q4' | 'int8' | 'fp16' | 'fp32';
  approxSizeMB: number;
  /** Extra options passed to tokenizer.apply_chat_template. */
  chatTemplateOptions?: Record<string, unknown>;
  notes?: string;
}

export const MODELS: ModelSpec[] = [
  {
    id: 'qwen3-0.6b',
    label: 'Qwen3-0.6B (q4f16, ~570MB)',
    repo: 'onnx-community/Qwen3-0.6B-ONNX',
    dtype: 'q4f16',
    approxSizeMB: 570,
    // Qwen3 has a thinking mode; disable it for this demo.
    chatTemplateOptions: { enable_thinking: false },
  },
  {
    id: 'qwen2.5-0.5b',
    label: 'Qwen2.5-0.5B-Instruct (q4f16, ~483MB)',
    repo: 'onnx-community/Qwen2.5-0.5B-Instruct',
    dtype: 'q4f16',
    approxSizeMB: 483,
  },
  {
    id: 'smollm2-360m',
    label: 'SmolLM2-360M-Instruct (q4f16, ~273MB)',
    repo: 'HuggingFaceTB/SmolLM2-360M-Instruct',
    dtype: 'q4f16',
    approxSizeMB: 273,
    notes: 'English-centric; lightest reasonable option.',
  },
];

export const DEFAULT_MODEL_ID = 'qwen3-0.6b';

/**
 * Sentence encoder for k-SemStamp. EmbeddingGemma is multilingual, so the
 * semantic clusters mean something in any language the demo model can write.
 *
 * Its ONNX graph already contains the pooling and the dense projection and
 * exposes `sentence_embedding` directly, so we call the model rather than the
 * feature-extraction pipeline - the pipeline would mean-pool the hidden states
 * and silently skip the projection, giving a different space.
 */
export const EMBEDDING_MODEL = {
  repo: 'onnx-community/embeddinggemma-300m-ONNX',
  dim: 768,
  approxSizeMB: 167,
  /**
   * EmbeddingGemma is trained with task prefixes. Clustering has its own, and
   * centroid computation and runtime assignment MUST use the same one or the
   * two embedding spaces do not line up.
   */
  clusteringPrefix: 'task: clustering | query: ',
  /**
   * Quantization is part of the embedding geometry, not just a size knob:
   * q4 and fp32 embeddings of one sentence sit at cosine ~0.98. The centroid
   * script pins the same value.
   */
  dtype: 'q4' as const,
};

export function getModelSpec(id: string): ModelSpec {
  const m = MODELS.find((m) => m.id === id);
  if (!m) throw new Error(`Unknown model id: ${id}`);
  return m;
}