cinny/src/app/hooks/useMessageTranslation.ts

import { useCallback, useEffect, useRef, useState } from 'react';
import { useAtom } from 'jotai';
import { settingsAtom } from '../state/settings';
import { useSetting } from '../state/hooks/settings';
import {
  translationCacheAtom,
  msgTranslationActiveAtomFamily,
  makeCacheKey,
  findTranslation,
  addTranslation,
} from '../state/translation';
import { chromeTranslationEngine as engine } from '../utils/translation/chromeEngine';
import { sameLanguage } from '../utils/translation/langUtils';

export type MsgTranslationStatus =
  | 'idle' // not translating (original shown)
  | 'detecting' // detecting source language
  | 'downloading' // downloading the on-device model (first use of a pair)
  | 'translating' // running the translation
  | 'done' // translation available
  | 'skipped' // already in the target language (nothing to do)
  | 'error'; // detection/translation failed

export type MsgTranslation = {
  /** Whether the on-device engine is usable at all in this browser. */
  supported: boolean;
  /** Whether the translated body should be shown in place of the original. */
  active: boolean;
  status: MsgTranslationStatus;
  translated?: string;
  /** Detected source language (normalized base code). */
  fromLang?: string;
  /** 0..1 while `status === 'downloading'`. */
  downloadProgress?: number;
  /** Toggle translation on/off for this message (satisfies the user gesture). */
  toggle: () => void;
};

/**
 * Per-message on-device translation. The message menu action and the inline
 * "Show original" chip both drive the same per-event `active` atom; whichever
 * component mounts the body runs the actual detect→translate flow (once) and
 * writes the result to the local cache. Nothing ever leaves the device.
 */
export const useMessageTranslation = (eventId: string, text: string): MsgTranslation => {
  const supported = engine.isSupported();
  const [targetLang] = useSetting(settingsAtom, 'translateTargetLang');
  const [autoTranslate] = useSetting(settingsAtom, 'autoTranslate');
  const [active, setActive] = useAtom(msgTranslationActiveAtomFamily(eventId));
  const [cache, setCache] = useAtom(translationCacheAtom);

  // The key includes a fingerprint of `text`, so an edit (same event id, new
  // body) misses the old entry and re-translates instead of showing stale text.
  const key = makeCacheKey(eventId, targetLang, text);
  const cached = findTranslation(cache, key);

  const [status, setStatus] = useState<MsgTranslationStatus>('idle');
  const [translated, setTranslated] = useState<string | undefined>(undefined);
  const [fromLang, setFromLang] = useState<string | undefined>(undefined);
  const [downloadProgress, setDownloadProgress] = useState<number | undefined>(undefined);

  const toggle = useCallback(() => setActive((a) => !a), [setActive]);

  // Adopt a cached translation immediately (scrollback / re-render never
  // re-translates).
  useEffect(() => {
    if (cached) {
      setTranslated(cached.translated);
      setFromLang(cached.fromLang);
      if (active) setStatus('done');
    }
  }, [cached, active]);

  // Auto-translate: when enabled, flip `active` on for foreign messages whose
  // model is already downloaded (no gesture/download needed). Messages needing
  // a first-time download keep the manual chip (which provides the gesture).
  // Store the text we last auto-tried so an edit (new text) re-triggers auto
  // detection instead of being suppressed by a one-shot flag.
  const autoTriedText = useRef<string | undefined>(undefined);
  useEffect(() => {
    if (!supported || !autoTranslate || active || cached || !eventId) return;
    if (!text.trim() || autoTriedText.current === text) return;
    autoTriedText.current = text;
    let alive = true;
    (async () => {
      const from = await engine.detect(text);
      if (!alive || !from || sameLanguage(from, targetLang)) return;
      const avail = await engine.availability(from, targetLang);
      if (!alive || avail !== 'available') return;
      setActive(true);
    })();
    return () => {
      alive = false;
    };
  }, [supported, autoTranslate, active, cached, text, targetLang, eventId, setActive]);

  // The detect→translate flow. Runs when translation is switched on and we have
  // no cached result. Guarded against overlapping runs / stale writes.
  const runIdRef = useRef(0);
  useEffect(() => {
    if (!active || cached || !supported) return undefined;
    if (!text.trim()) {
      setStatus('skipped');
      return undefined;
    }

    const runId = runIdRef.current + 1;
    runIdRef.current = runId;
    const stale = () => runId !== runIdRef.current;

    (async () => {
      try {
        setStatus('detecting');
        setDownloadProgress(undefined);
        const from = await engine.detect(text);
        if (stale()) return;
        if (!from || sameLanguage(from, targetLang)) {
          setStatus('skipped');
          setActive(false);
          return;
        }
        setFromLang(from);

        const avail = await engine.availability(from, targetLang);
        if (stale()) return;
        if (avail === 'unavailable') {
          setStatus('error');
          return;
        }
        setStatus(avail === 'available' ? 'translating' : 'downloading');

        const result = await engine.translate(text, from, targetLang, (p) => {
          if (!stale()) setDownloadProgress(p);
        });
        if (stale()) return;

        setTranslated(result);
        setDownloadProgress(undefined);
        setStatus('done');
        setCache((prev) => addTranslation(prev, { key, fromLang: from, translated: result }));
      } catch {
        if (!stale()) setStatus('error');
      }
    })();

    return () => {
      // Invalidate this run's writes if inputs change / unmount.
      runIdRef.current += 1;
    };
  }, [active, cached, supported, text, targetLang, key, setActive, setCache]);

  // When switched off, drop back to idle (keep the cached translation around so
  // toggling back on is instant).
  useEffect(() => {
    if (!active && status !== 'idle') setStatus('idle');
  }, [active, status]);

  return {
    supported,
    active,
    status: active ? status : 'idle',
    translated: cached?.translated ?? translated,
    fromLang: cached?.fromLang ?? fromLang,
    downloadProgress,
    toggle,
  };
};

/** Lightweight helper for the message menu — is the translate action offerable? */
export const isTranslationSupported = (): boolean => engine.isSupported();