/**
 * Opus-MT (MarianMT) through Transformers.js — the model half of transcript
 * translation (D-WS4b). Shared by the browser translate worker and the
 * onnxruntime-node measurement script / real-model test, so all three run
 * the exact same pipeline options.
 */

import { pipeline, type ProgressInfo, type TranslationPipeline } from "@huggingface/transformers";

/** Hard input limit of every Opus-MT checkpoint (max_position_embeddings). */
export const OPUS_MT_MAX_POSITIONS = 512;

export interface OpusMtTranslator {
  /** One translated string per piece, same order. */
  translateBatch(pieces: string[]): Promise<string[]>;
  /** Real SentencePiece token count (incl. </s>) for the splitter's budget. */
  countTokens(text: string): number;
  dispose(): Promise<void>;
}

export interface LoadOpusMtOptions {
  /** "wasm" in the browser; omit under onnxruntime-node (defaults to CPU). */
  device?: "wasm";
  progressCallback?: (info: ProgressInfo) => void;
}

export async function loadOpusMt(
  modelId: string,
  options: LoadOpusMtOptions = {},
): Promise<OpusMtTranslator> {
  const translator = (await pipeline("translation", modelId, {
    // q8: 113 MB per direction (encoder + merged decoder + tokenizer); the
    // fp32 export is ~4x larger for no audible gain on transcripts.
    dtype: "q8",
    ...(options.device && { device: options.device }),
    // D-014: the bundled onnxruntime-web's level-2 QDQ transform rejects
    // 8-bit decoders with a tied embedding weight — Marian ties them too.
    session_options: { graphOptimizationLevel: "basic" as const },
    ...(options.progressCallback && { progress_callback: options.progressCallback }),
  })) as TranslationPipeline;

  const countTokens = (text: string) => translator.tokenizer.encode(text).length;

  return {
    countTokens,
    async translateBatch(pieces) {
      if (pieces.length === 0) return [];
      // Cap output length relative to the longest input: a translation is
      // rarely more than ~2x its source in tokens, and the cap stops a rare
      // repetition loop from running to 512 tokens.
      const longest = Math.max(...pieces.map(countTokens));
      const output = (await translator(pieces, {
        max_new_tokens: Math.min(OPUS_MT_MAX_POSITIONS, 3 * longest + 16),
      })) as { translation_text: string }[];
      return output.map((entry) => entry.translation_text);
    },
    dispose: () => translator.dispose(),
  };
}