/**
 * Text translation of a finished transcript (D-WS4b) — the pure half.
 *
 * Opus-MT (MarianMT) runs in `src/worker/translate.worker.ts`; everything
 * here is planning that needs no model: which checkpoint serves a language
 * pair, how a transcript becomes translation units that keep their
 * timestamps, and how a unit is cut into sentence pieces under the model's
 * input budget. No browser globals, so `bun test` covers it.
 *
 * Why per-sentence pieces: Opus-MT is trained on sentence pairs and, given
 * several sentences at once, routinely drops the later ones (measured
 * 2026-10-29: "…I live in Copenhagen. How are you?" came back as one
 * sentence). Its encoder also tops out at 512 positions.
 */

import type { TranscriptChunk } from "../worker/protocol";

export type TranslationLanguage = "da" | "en";

export interface TranslationModel {
  /** Hugging Face Hub id of the Transformers.js ONNX export. */
  id: string;
  source: TranslationLanguage;
  target: TranslationLanguage;
  /** Direction as shown in the picker, e.g. "Danish → English". */
  label: string;
  /** Plain name of the target language, e.g. "English". */
  targetLabel: string;
  /**
   * One-time download at the q8 dtype the worker loads: encoder + merged
   * decoder + tokenizer/configs, summed from the Hub file sizes.
   */
  approxBytes: number;
}

/**
 * Xenova's ONNX exports of Helsinki-NLP's Opus-MT checkpoints (Apache-2.0
 * upstream). Verified to load and translate under Transformers.js 4.2.0.
 */
export const TRANSLATION_MODELS: readonly TranslationModel[] = [
  {
    id: "Xenova/opus-mt-da-en",
    source: "da",
    target: "en",
    label: "Danish → English",
    targetLabel: "English",
    approxBytes: 113_000_000,
  },
  {
    id: "Xenova/opus-mt-en-da",
    source: "en",
    target: "da",
    label: "English → Danish",
    targetLabel: "Danish",
    approxBytes: 113_000_000,
  },
];

export const TRANSLATION_MODEL_IDS: readonly string[] = TRANSLATION_MODELS.map((m) => m.id);

export function getTranslationModel(
  source: TranslationLanguage,
  target: TranslationLanguage,
): TranslationModel | undefined {
  return TRANSLATION_MODELS.find((m) => m.source === source && m.target === target);
}

/** Lookup by Hub id (storage panel names, worker request validation). */
export function getTranslationModelById(id: string): TranslationModel | undefined {
  return TRANSLATION_MODELS.find((m) => m.id === id);
}

// ---------------------------------------------------------------------------
// Units: transcript → sentence-ending groups of chunks with their time span
// ---------------------------------------------------------------------------

/** A stretch of transcript translated as one cue. */
export interface TranslationUnit {
  text: string;
  /** Merged [start, end] of the source chunks; null for untimed text. */
  timestamp: [number, number | null] | null;
}

const SENTENCE_END = /[.!?…]["'”’)\]]*$/;

/** Default cap on chunks merged into one unit, so an unpunctuated transcript
 * still yields subtitle-sized cues rather than one giant one. */
export const MAX_CHUNKS_PER_UNIT = 8;

/**
 * Group a transcript into translation units. Whisper's segments often end
 * mid-sentence; translating fragments hurts quality, so consecutive chunks
 * are merged until one ends a sentence (or `maxChunksPerUnit` is reached)
 * and the unit takes the merged time span — `.srt` keeps working, with
 * slightly longer cues. Without chunks (nothing timestamped), each non-empty
 * line is an untimed unit.
 */
export function buildTranslationUnits(
  result: { text: string; chunks: TranscriptChunk[] },
  maxChunksPerUnit: number = MAX_CHUNKS_PER_UNIT,
): TranslationUnit[] {
  const chunks = result.chunks.filter((chunk) => chunk.text.trim() !== "");
  if (chunks.length === 0) {
    return result.text
      .split(/\n+/)
      .map((line) => line.trim())
      .filter((line) => line !== "")
      .map((text) => ({ text, timestamp: null }));
  }
  const units: TranslationUnit[] = [];
  let group: TranscriptChunk[] = [];
  const flush = () => {
    const first = group[0];
    const last = group[group.length - 1];
    if (!first || !last) return;
    units.push({
      text: group.map((chunk) => chunk.text.trim()).join(" "),
      timestamp: [first.timestamp[0], last.timestamp[1]],
    });
    group = [];
  };
  for (const chunk of chunks) {
    group.push(chunk);
    if (SENTENCE_END.test(chunk.text.trim()) || group.length >= maxChunksPerUnit) flush();
  }
  flush();
  return units;
}

// ---------------------------------------------------------------------------
// Pieces: unit text → sentences under the model's token budget
// ---------------------------------------------------------------------------

/**
 * Cheap upper-bound token estimate for Opus-MT's SentencePiece vocab, used
 * when the real tokenizer isn't at hand (tests, planning). Measured ~0.34
 * tokens/char on en_alien and ~0.26 on da_festuge with the real tokenizers;
 * 0.5/char plus the end-of-sentence token stays above both.
 */
export function estimateTokens(text: string): number {
  return text.length === 0 ? 0 : Math.ceil(text.length / 2) + 1;
}

/**
 * Default per-piece budget, well inside Opus-MT's 512 positions: quality
 * falls off on very long inputs long before the hard limit, and normal
 * sentences are 10–60 tokens, so this only bites on unpunctuated run-ons.
 */
export const DEFAULT_MAX_PIECE_TOKENS = 200;

export interface SplitOptions {
  maxTokens?: number;
  /** Token counter — the worker passes the real tokenizer. */
  countTokens?: (text: string) => number;
}

/**
 * Title abbreviations (en + da) and single-letter initials ("J.") — the
 * sentence segmenter's false breaks. Kept to an explicit list: a broad
 * "short capitalised word" rule would glue real one-word sentences
 * ("Yes.", "Okay.") onto the next one.
 */
const ABBREVIATION = /^(?:Mr|Mrs|Ms|Dr|Prof|St|Jr|Sr|Hr|Fru|Nr|\p{Lu})\.$/u;

function sentences(text: string): string[] {
  const out: string[] = [];
  let carry = "";
  for (const { segment } of new Intl.Segmenter("en", { granularity: "sentence" }).segment(text)) {
    const trimmed = segment.trim();
    if (trimmed === "") continue;
    if (ABBREVIATION.test(trimmed)) {
      carry = carry === "" ? trimmed : `${carry} ${trimmed}`;
      continue;
    }
    out.push(carry === "" ? trimmed : `${carry} ${trimmed}`);
    carry = "";
  }
  if (carry !== "") out.push(carry);
  return out;
}

/** Greedily pack `parts` (in order) into space-joined pieces under budget. */
function pack(parts: string[], maxTokens: number, count: (text: string) => number): string[] {
  const pieces: string[] = [];
  let current = "";
  for (const part of parts) {
    const candidate = current === "" ? part : `${current} ${part}`;
    if (current !== "" && count(candidate) > maxTokens) {
      pieces.push(current);
      current = part;
    } else {
      current = candidate;
    }
  }
  if (current !== "") pieces.push(current);
  return pieces;
}

/** Split one over-budget sentence: commas first, then words. A single word
 * longer than the budget is kept whole (nothing is ever dropped). */
function wrap(sentence: string, maxTokens: number, count: (text: string) => number): string[] {
  const clauses = sentence.split(/(?<=,)\s+/);
  const parts = clauses.flatMap((clause) =>
    count(clause) <= maxTokens ? [clause] : pack(clause.split(/\s+/), maxTokens, count),
  );
  return pack(parts, maxTokens, count);
}

/**
 * Cut a unit's text into the pieces actually fed to Opus-MT: one per
 * sentence, an over-budget sentence wrapped on commas then words. Order is
 * preserved and every word survives — joining the pieces with spaces gives
 * the input back up to whitespace.
 */
export function splitForTranslation(text: string, options: SplitOptions = {}): string[] {
  const maxTokens = options.maxTokens ?? DEFAULT_MAX_PIECE_TOKENS;
  const count = options.countTokens ?? estimateTokens;
  return sentences(text).flatMap((sentence) =>
    count(sentence) <= maxTokens ? [sentence] : wrap(sentence, maxTokens, count),
  );
}

// ---------------------------------------------------------------------------
// The translation loop (model injected: worker, benchmark and tests share it)
// ---------------------------------------------------------------------------

export interface TranslateUnitsDeps extends SplitOptions {
  /** Translate a batch of pieces; must answer one string per piece, in order. */
  translateBatch: (pieces: string[]) => Promise<string[]>;
  /** Pieces per model call (padding makes larger batches cheaper per piece). */
  batchSize?: number;
  /** Called with (done, total) pieces: once at 0, then after every batch. */
  onProgress?: (done: number, total: number) => void;
}

/**
 * One piece per model call. Batching would be ~4x faster, but Transformers.js
 * 4.2's batched Marian generation is broken: every row of a padded batch
 * comes back as row 0's translation (measured 2026-10-29 on
 * onnxruntime-node, q8 and fp32, any graph optimisation level — D-WS4b).
 * Re-measure with `scripts/bench-translate.ts` after a Transformers.js
 * upgrade before raising this.
 */
export const DEFAULT_TRANSLATE_BATCH_SIZE = 1;

/**
 * Translate unit texts: split each into sentence pieces, translate all
 * pieces in order-preserving batches, rejoin each unit's pieces with spaces.
 * Returns one translated string per input unit.
 */
export async function translateUnits(
  units: readonly string[],
  deps: TranslateUnitsDeps,
): Promise<string[]> {
  const batchSize = Math.max(1, deps.batchSize ?? DEFAULT_TRANSLATE_BATCH_SIZE);
  const perUnit = units.map((unit) => splitForTranslation(unit, deps));
  const pieces = perUnit.flat();
  const translated: string[] = [];
  deps.onProgress?.(0, pieces.length);
  for (let start = 0; start < pieces.length; start += batchSize) {
    const batch = pieces.slice(start, start + batchSize);
    const out = await deps.translateBatch(batch);
    if (out.length !== batch.length) {
      throw new Error(`Translation returned ${out.length} sentences for ${batch.length}.`);
    }
    translated.push(...out.map((text) => text.trim()));
    deps.onProgress?.(translated.length, pieces.length);
  }
  let next = 0;
  return perUnit.map((unitPieces) => {
    const joined = translated.slice(next, next + unitPieces.length).join(" ");
    next += unitPieces.length;
    return joined;
  });
}

// ---------------------------------------------------------------------------
// Assembly
// ---------------------------------------------------------------------------

/**
 * Pair translated unit texts back with their units' time spans, in the
 * `{ text, chunks }` shape the transcript view and exports already take.
 */
export function assembleTranslation(
  units: readonly TranslationUnit[],
  translated: readonly string[],
): { text: string; chunks: TranscriptChunk[] } {
  if (units.length !== translated.length) {
    throw new Error(
      `Translation returned ${translated.length} parts for ${units.length} — refusing to misalign the transcript.`,
    );
  }
  const texts = translated.map((text) => text.trim());
  const chunks: TranscriptChunk[] = [];
  units.forEach((unit, index) => {
    if (unit.timestamp) chunks.push({ text: texts[index] ?? "", timestamp: unit.timestamp });
  });
  const untimed = units.every((unit) => unit.timestamp === null);
  return { text: texts.join(untimed ? "\n" : " "), chunks };
}

// ---------------------------------------------------------------------------
// Source language guess (default for the direction picker)
// ---------------------------------------------------------------------------

const DANISH_WORDS = new Set([
  "og", "ikke", "det", "jeg", "vi", "på", "til", "med", "som", "har", "af", "der",
  "kan", "skal", "vil", "også", "men", "hvad", "hvordan", "være", "noget",
]);
const ENGLISH_WORDS = new Set([
  "the", "and", "is", "to", "of", "that", "you", "it", "with", "have", "this", "not",
  "what", "how", "be", "was", "are",
]);

/**
 * Best guess at which language a transcript is in, for the direction
 * picker's default (the user can always change it). An explicit Whisper
 * language wins; otherwise Danish letters and function words vs English
 * ones decide, defaulting to English.
 */
export function guessSourceLanguage(text: string, whisperLanguage?: string): TranslationLanguage {
  if (whisperLanguage === "da" || whisperLanguage === "en") return whisperLanguage;
  const lower = text.toLowerCase();
  let danish = (lower.match(/[æøå]/g)?.length ?? 0) * 2;
  let english = 0;
  for (const word of lower.match(/\p{L}+/gu) ?? []) {
    if (DANISH_WORDS.has(word)) danish += 1;
    if (ENGLISH_WORDS.has(word)) english += 1;
  }
  return danish > english ? "da" : "en";
}