SkateSpots/src/lib/translation.ts

/**
 * On-device translation client.
 *
 * The actual model runs in `workers/translate.worker.ts`. This module only
 * exchanges `postMessage` with it; it never fetches message text and keeps
 * translations in memory only (no IndexedDB/localStorage).
 */
import type { Language } from '@/contexts/language';

export interface TranslationProgress {
  progress?: number;
  status?: string;
}

export interface TranslateParams {
  id: string;
  text: string;
  source: Language;
  target: Language;
  onProgress?: (progress: TranslationProgress) => void;
}

type WorkerRequest = {
  type: 'translate';
  id: number;
  text: string;
  source: Language;
  target: Language;
};

type WorkerResponse =
  | { type: 'progress'; id: number; status?: string; progress?: number }
  | { type: 'result'; id: number; text: string }
  | { type: 'error'; id: number; code: string };

interface PendingRequest {
  key: string;
  resolve: (text: string) => void;
  reject: (error: Error) => void;
  onProgress?: (progress: TranslationProgress) => void;
  timer: ReturnType<typeof setTimeout>;
}

const MODEL_TIMEOUT_MS = 300_000; // first use downloads a ~110 MB model
const memoryCache = new Map<string, string>();
const pending = new Map<number, PendingRequest>();
let worker: Worker | null = null;
let sequence = 0;

/** The pilot supports English ↔ Spanish. */
export function isTranslationPairSupported(
  source: Language | null | undefined,
  target: Language,
): boolean {
  if (!source) return false;
  return (source === 'en' && target === 'es') || (source === 'es' && target === 'en');
}

function rejectAll(error: Error): void {
  for (const request of pending.values()) {
    clearTimeout(request.timer);
    request.reject(error);
  }
  pending.clear();
}

function getWorker(): Worker {
  if (worker) return worker;

  const created = new Worker(
    new URL('../workers/translate.worker.ts', import.meta.url),
    { type: 'module' },
  );

  created.onmessage = (event: MessageEvent<WorkerResponse>) => {
    const data = event.data;
    if (!data || typeof data.id !== 'number') return;
    const request = pending.get(data.id);
    if (!request) return;

    if (data.type === 'progress') {
      request.onProgress?.({ progress: data.progress, status: data.status });
      return;
    }

    pending.delete(data.id);
    clearTimeout(request.timer);
    if (data.type === 'result') {
      memoryCache.set(request.key, data.text);
      request.resolve(data.text);
    } else {
      request.reject(new Error(data.code || 'FAILED'));
    }
  };

  created.onerror = () => {
    rejectAll(new Error('WORKER'));
    worker?.terminate();
    worker = null;
  };

  worker = created;
  return created;
}

/** Translate a message on-device. The text is never sent over the network. */
export async function translateText(params: TranslateParams): Promise<string> {
  if (params.source === params.target) return params.text;
  if (!isTranslationPairSupported(params.source, params.target)) {
    throw new Error('UNSUPPORTED');
  }

  const key = `${params.id}:${params.target}`;
  const cached = memoryCache.get(key);
  if (cached) return cached;

  const activeWorker = getWorker();
  const id = ++sequence;

  return new Promise<string>((resolve, reject) => {
    const timer = setTimeout(() => {
      pending.delete(id);
      reject(new Error('TIMEOUT'));
    }, MODEL_TIMEOUT_MS);

    pending.set(id, {
      key,
      resolve,
      reject,
      onProgress: params.onProgress,
      timer,
    });

    const request: WorkerRequest = {
      type: 'translate',
      id,
      text: params.text,
      source: params.source,
      target: params.target,
    };
    activeWorker.postMessage(request);
  });
}

/** Clear the in-memory translation cache (e.g. on logout). */
export function clearTranslationMemory(): void {
  memoryCache.clear();
  rejectAll(new Error('CLEARED'));
  if (worker) {
    worker.terminate();
    worker = null;
  }
}