// Auto-captions using @xenova/transformers' Whisper.
//
// Pipeline:
//   MediaAsset.audioBuffer (or File) → 16kHz mono Float32 PCM
//   → Whisper AutomaticSpeechRecognition pipeline (ONNX, runs in Web Worker)
//   → chunks [{ text, timestamp: [start, end] }]
//   → hallucination filter → dedupe → SubtitleCue[] merged into editor.
//
// Why whisper-tiny as default? It's the smallest multilingual model that
// supports Thai (~40 MB quantized). Larger models (base/small) are 2-5× more
// accurate but take 3-10× longer per minute. We expose model size as a
// setting so users can pick.
//
// Output cleaning (the important bit):
//   Whisper-tiny frequently **hallucinates** when the input is music, noise,
//   or low-energy audio — producing Thai text with repeated vowel marks like
//   "อนุๆใจปล๊ก๊กเลย" or "ที่ส่งด๊อด๊า่าลืมไป๊ก๊เลย". We:
//     1. Drop cues that are pure Thai diacritic soup (no consonants)
//     2. Drop cues where the same syllable repeats 3+ times in a row
//     3. Drop cues whose romanized length vs Thai-char ratio is suspicious
//     4. Collapse repeated vowel-mark runs ("๊ก๊ก๊" → "๊ก")
//     5. Strip leading/trailing punctuation garbage
//
// Caching: we use our own IndexedDB-backed cache (see ./model-cache) and
// intercept fetches so transformers.js doesn't depend on the Cache API.
// This works in headless browsers / restricted environments where Cache
// API isn't available.

import type { SubtitleCue } from '@/types';
import { prewarmWhisperModel, installCacheInterception, ensureWhisperModel } from './model-cache';

export type WhisperModelSize = 'tiny' | 'base';

export interface CaptionOptions {
  language?: 'th' | 'en' | 'auto';
  /** Max chunk length in seconds. Whisper handles up to 30s well. */
  chunkSeconds?: number;
  /** Model size — tiny (default, ~40MB) or base (~75MB, more accurate). */
  modelSize?: WhisperModelSize;
}

export interface CaptionProgress {
  phase: 'init' | 'decode' | 'transcribe' | 'merge';
  message: string;
  pct: number;
}

export type CaptionProgressCallback = (p: CaptionProgress) => void;

// Cache pip promise per model id so user can switch models without rebuilding.
const pipelinePromises = new Map<string, Promise<unknown>>();

function modelIdFor(size: WhisperModelSize): string {
  return size === 'base' ? 'Xenova/whisper-base' : 'Xenova/whisper-tiny';
}

async function getPipeline(
  modelSize: WhisperModelSize,
  onProgress?: CaptionProgressCallback,
): Promise<unknown> {
  const modelId = modelIdFor(modelSize);
  const cached = pipelinePromises.get(modelId);
  if (cached) return cached;

  const promise = (async () => {
    onProgress?.({ phase: 'init', message: 'กำลังเตรียม Whisper model...', pct: 5 });
    const expectedFiles = modelSize === 'base'
      ? 8  // base has more config files
      : 6;
    await ensureWhisperModel(modelId, (file, idx) => {
      onProgress?.({
        phase: 'init',
        message: `ดาวน์โหลด ${file} (${idx}/${expectedFiles})`,
        pct: 5 + Math.round((idx / expectedFiles) * 30),
      });
    });
    installCacheInterception();
    const { pipeline, env } = await import('@xenova/transformers');
    env.allowLocalModels = false;
    env.useBrowserCache = typeof caches !== 'undefined';
    onProgress?.({ phase: 'init', message: 'กำลังโหลด Whisper pipeline...', pct: 38 });
    const p = await pipeline('automatic-speech-recognition', modelId, {
      quantized: true,
    } as Parameters<typeof pipeline>[2]);
    onProgress?.({ phase: 'init', message: 'พร้อมถอดเสียง', pct: 42 });
    return p;
  })();

  pipelinePromises.set(modelId, promise);
  return promise;
}

/** Convert an AudioBuffer or raw Float32Array to 16kHz mono PCM. */
export function toMono16k(
  source: AudioBuffer | Float32Array,
  sourceSampleRate = 44100,
): Float32Array {
  const targetRate = 16000;
  let mono: Float32Array;
  let srcRate = sourceSampleRate;

  if (source instanceof Float32Array) {
    mono = source;
  } else {
    // Average all channels → mono.
    const ch = source.numberOfChannels;
    const len = source.length;
    mono = new Float32Array(len);
    for (let c = 0; c < ch; c++) {
      const data = source.getChannelData(c);
      for (let i = 0; i < len; i++) mono[i] += data[i];
    }
    for (let i = 0; i < len; i++) mono[i] /= ch;
    srcRate = source.sampleRate;
  }

  // Resample to 16kHz with linear interpolation. Cheap and good enough for
  // ASR.
  if (srcRate === targetRate) return mono;
  const ratio = srcRate / targetRate;
  const outLen = Math.floor(mono.length / ratio);
  const out = new Float32Array(outLen);
  for (let i = 0; i < outLen; i++) {
    const idx = i * ratio;
    const i0 = Math.floor(idx);
    const i1 = Math.min(i0 + 1, mono.length - 1);
    const t = idx - i0;
    out[i] = mono[i0] * (1 - t) + mono[i1] * t;
  }
  return out;
}

interface WhisperChunk {
  text: string;
  timestamp: [number, number] | null;
}

/** Run caption inference on an AudioBuffer. */
export async function generateCaptions(
  audio: AudioBuffer,
  options: CaptionOptions = {},
  onProgress?: CaptionProgressCallback,
): Promise<SubtitleCue[]> {
  const { language = 'auto', chunkSeconds = 25, modelSize = 'tiny' } = options;

  onProgress?.({ phase: 'decode', message: 'เตรียม PCM 16kHz...', pct: 12 });
  const pcm = toMono16k(audio, audio.sampleRate);
  const totalSeconds = pcm.length / 16000;

  const transcriber = (await getPipeline(modelSize, onProgress)) as (
    input: Float32Array,
    config: Record<string, unknown>,
  ) => Promise<{ chunks?: WhisperChunk[]; text?: string }>;

  // We process the audio in overlapping windows so each chunk has its own
  // timestamps. Whisper's `return_timestamps: true` returns offsets in the
  // input buffer; we adjust them by the chunk's start.
  const chunkSamples = chunkSeconds * 16000;
  const overlapSamples = 2 * 16000; // 2s overlap to avoid cutting words
  const allChunks: { start: number; end: number; text: string }[] = [];

  let offsetSamples = 0;
  let iter = 0;
  while (offsetSamples < pcm.length) {
    const end = Math.min(offsetSamples + chunkSamples, pcm.length);
    const slice = pcm.subarray(offsetSamples, end);
    const startSec = offsetSamples / 16000;
    const endSec = end / 16000;
    onProgress?.({
      phase: 'transcribe',
      message: `ถอดเสียง ${startSec.toFixed(1)}s – ${endSec.toFixed(1)}s / ${totalSeconds.toFixed(1)}s`,
      pct: 45 + Math.min(45, (end / pcm.length) * 45),
    });

    const result = await transcriber(slice, {
      language: language === 'auto' ? null : language,
      task: 'transcribe',
      return_timestamps: 'word' as const,
      chunk_length_s: chunkSeconds,
      stride_length_s: 2,
    });

    if (result.chunks && result.chunks.length > 0) {
      for (const c of result.chunks) {
        const [localStart, localEnd] = c.timestamp ?? [0, 0];
        allChunks.push({
          start: startSec + (localStart ?? 0),
          end: startSec + (localEnd ?? 0),
          text: cleanText(c.text),
        });
      }
    } else if (result.text) {
      // Fallback: no per-word timestamps → split evenly.
      const segText = cleanText(result.text);
      if (segText) {
        allChunks.push({
          start: startSec,
          end: endSec,
          text: segText,
        });
      }
    }

    offsetSamples += chunkSamples - overlapSamples;
    iter++;
    if (iter > 200) break; // safety net
  }

  onProgress?.({ phase: 'merge', message: 'กรอง hallucination...', pct: 92 });
  const filtered = allChunks.filter((c) => !isHallucination(c.text, c.end - c.start));
  const merged = mergeAdjacent(filtered, 1.5);
  // If everything got filtered (no real speech detected), return empty array.
  onProgress?.({ phase: 'merge', message: 'เสร็จสิ้น', pct: 100 });
  return merged;
}

/** Strip Whisper special tokens and collapse whitespace. */
function cleanText(text: string): string {
  return text
    .replace(/<\|.*?\|>/g, '')
    .replace(/\s+/g, ' ')
    .trim();
}

/**
 * Detect Whisper hallucination — text that's clearly noise rather than
 * speech. Whisper-tiny emits this on music / silent / non-speech audio.
 *
 * The hard cases are Thai hallucinations where the model randomly samples
 * Thai Unicode characters — the result looks superficially like Thai
 * (vowel + consonant ratios are reasonable, character set is Thai) but is
 * pure nonsense. We catch these with a **char-rate** check: real Thai
 * speech produces ≤ ~15 chars/second; hallucinated Thai on a 2-3 second
 * chunk is often 20+ chars/second because the model "speeds through" the
 * fake text.
 */
export function isHallucination(text: string, durationSec = 0): boolean {
  if (!text || text.length < 2) return true;
  // Drop common "thank you for watching" hallucinations from Whisper.
  const lower = text.toLowerCase();
  if (
    /thank you for watching|subscrib|like and|please like|subtitle by|translated by/i.test(lower)
  ) return true;
  // Drop strings made mostly of punctuation / spaces / numbers.
  const stripped = text.replace(/[\s\d.,!?()"'\-–—…:;]/g, '');
  if (stripped.length < 2) return true;
  // Runaway diacritic soup: 3+ diacritic/vowel marks in a row.
  if (/[\u0E48-\u0E4C\u0E31\u0E34-\u0E3A]{3,}/.test(text)) return true;
  // Same single character 5+ times in a row.
  if (/(.)\1{4,}/.test(stripped)) return true;
  // Same 1-3 char syllable repeated 4+ times in a row.
  if (/(.{1,3})\1{3,}/.test(stripped)) return true;
  // Pure English with > 3 identical chars in a row.
  if (/([a-zA-Z])\1{4,}/.test(text)) return true;
  // Char-rate check (the most reliable Thai hallucination detector).
  // Real Thai speech: ~5-14 chars/second. Hallucinated chunks often 20+.
  // We use 18 chars/sec as the threshold — catches Whisper's "speed
  // through garbage" hallucinations while leaving room for fast speakers.
  if (durationSec > 0.5) {
    const charsPerSec = stripped.length / durationSec;
    if (charsPerSec > 18) return true;
  }
  // Very short cues with very few Thai characters are usually noise
  // (Whisper emits things like "อณณ สู่" on short low-energy chunks).
  if (stripped.length < 4) return true;
  // Repeated-pair signature for Thai-heavy text.
  const thaiConsonants = (stripped.match(/[\u0E01-\u0E2E]/g) || []).length;
  const diacritics = (stripped.match(/[\u0E31\u0E34-\u0E3A\u0E47-\u0E4E]/g) || []).length;
  const vowels = (stripped.match(/[\u0E30-\u0E3A\u0E40-\u0E44]/g) || []).length;
  if (thaiConsonants >= 6) {
    // Vowel marks must be ≥ 20% of consonants (real Thai has vowel in every syllable).
    if (vowels / thaiConsonants < 0.20) return true;
    // More diacritics than consonants = nonsense Thai.
    if (diacritics / thaiConsonants > 1.3) return true;
  }
  // Repeated 2-char pair signature.
  if (stripped.length >= 8) {
    const pairs = new Map<string, number>();
    for (let i = 0; i < stripped.length - 1; i++) {
      const pair = stripped.slice(i, i + 2);
      pairs.set(pair, (pairs.get(pair) ?? 0) + 1);
    }
    let maxRepeat = 0;
    let repeatPairs = 0;
    for (const [, count] of pairs) {
      if (count >= 3) repeatPairs++;
      if (count > maxRepeat) maxRepeat = count;
    }
    // If any single 2-char pair appears 4+ times in a row of 8+, hallucination.
    if (maxRepeat >= 4) return true;
    // If 30%+ of distinct pairs are repeated 3+ times → cluster of repeats.
    if (pairs.size > 0 && repeatPairs / pairs.size > 0.3) return true;
  }
  return false;
}

/**
 * Collapse obvious repeated-run garbage that slips past isHallucination:
 * "ปล๊ก๊ก๊" → "ปล๊ก", "ดดด" → "ด". Applied per-character with limit
 * of 2 consecutive repeats.
 */
function collapseRepeats(text: string): string {
  return text.replace(/(.)\1{2,}/g, '$1$1');
}

/** Merge adjacent chunks whose texts belong to the same sentence. */
function mergeAdjacent(
  chunks: { start: number; end: number; text: string }[],
  maxGap: number,
): SubtitleCue[] {
  const merged: { start: number; end: number; text: string }[] = [];
  for (const c of chunks) {
    const text = collapseRepeats(c.text);
    if (!text) continue;
    const last = merged[merged.length - 1];
    if (last && c.start - last.end <= maxGap && last.text.length < 80) {
      last.text = (last.text + ' ' + text).trim();
      last.end = c.end;
    } else {
      merged.push({ start: c.start, end: c.end, text });
    }
  }
  return merged.map((m, i) => ({
    id: i + 1,
    start: m.start,
    end: m.end,
    text: capitalise(m.text),
  }));
}

function capitalise(s: string): string {
  if (!s) return s;
  return s.charAt(0).toLocaleUpperCase() + s.slice(1);
}