background.ts

/**
 * entrypoints/background.ts
 *
 * Persistent background service worker (Manifest V3).
 * Owns the single WebSocket connection to the Fact Engine backend.
 * Receives analysis batches from content scripts → forwards to backend.
 * Receives analysis results from backend → routes to correct tab's content script.
 *
 * Reconnect strategy: exponential backoff starting at 1s, max 30s cap.
 * The extension icon badge reflects WS state.
 */

import { defineBackground } from 'wxt/sandbox';

export default defineBackground({
  type: 'module',
  persistent: true,

  main() {
    const WS_URL = (import.meta.env.VITE_WS_URL || 'ws://localhost:7860');
    const CLIENT_ID = getOrCreateClientId();

    let ws: WebSocket | null = null;
    let reconnectDelay = 1000;
    let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
    let isConnected = false;

    // Map: tab content scripts waiting for results
    const pendingTabs = new Map<string, number>(); // claimHash → tabId

    // ── Badge helpers ─────────────────────────────────────────────────────
    function setBadge(status: 'connected' | 'reconnecting' | 'offline') {
      const colors = {
        connected: '#22c55e',
        reconnecting: '#eab308',
        offline: '#6b7280',
      };
      const labels = { connected: '●', reconnecting: '↻', offline: '✕' };
      chrome.action.setBadgeText({ text: labels[status] });
      chrome.action.setBadgeBackgroundColor({ color: colors[status] });
      // Notify all content scripts of status change
      chrome.tabs.query({}, (tabs) => {
        tabs.forEach((tab) => {
          if (tab.id) {
            chrome.tabs.sendMessage(tab.id, { type: 'ws_status', status }).catch(() => {});
          }
        });
      });
    }

    // ── WebSocket connection ──────────────────────────────────────────────
    function connect() {
      if (ws && ws.readyState === WebSocket.OPEN) return;

      const url = `${WS_URL}/ws/${CLIENT_ID}`;
      console.log('[FactEngine] Connecting to', url);
      setBadge('reconnecting');

      try {
        ws = new WebSocket(url);
      } catch (err) {
        console.error('[FactEngine] WebSocket construction failed', err);
        scheduleReconnect();
        return;
      }

      ws.onopen = () => {
        console.log('[FactEngine] WebSocket connected');
        isConnected = true;
        reconnectDelay = 1000; // reset backoff on successful connection
        setBadge('connected');
      };

      ws.onmessage = (event: MessageEvent) => {
        try {
          const msg = JSON.parse(event.data as string);
          handleBackendMessage(msg);
        } catch (err) {
          console.error('[FactEngine] Failed to parse backend message', err);
        }
      };

      ws.onerror = (err) => {
        console.error('[FactEngine] WebSocket error', err);
      };

      ws.onclose = () => {
        console.log('[FactEngine] WebSocket closed');
        isConnected = false;
        ws = null;
        setBadge('offline');
        scheduleReconnect();
      };
    }

    function scheduleReconnect() {
      if (reconnectTimer) clearTimeout(reconnectTimer);
      const jitter = Math.random() * 500;
      reconnectTimer = setTimeout(() => {
        connect();
        reconnectDelay = Math.min(reconnectDelay * 2, 30_000); // cap at 30s
      }, reconnectDelay + jitter);
    }

    function sendToBackend(payload: object) {
      if (ws && ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify(payload));
        return true;
      }
      return false;
    }

    // ── Handle messages from backend ──────────────────────────────────────
    function handleBackendMessage(msg: Record<string, unknown>) {
      if (msg.type === 'analysis_batch') {
        const results = msg.results as Array<Record<string, unknown>>;
        // Route each result to all active content script tabs
        chrome.tabs.query({ active: true }, (tabs) => {
          tabs.forEach((tab) => {
            if (tab.id) {
              chrome.tabs.sendMessage(tab.id, {
                type: 'analysis_results',
                results,
              }).catch(() => {/* tab may not have content script */});
            }
          });
        });
        // Also broadcast to all tabs (for multi-tab usage)
        chrome.tabs.query({}, (tabs) => {
          tabs.forEach((tab) => {
            if (tab.id) {
              chrome.tabs.sendMessage(tab.id, {
                type: 'analysis_results',
                results,
              }).catch(() => {});
            }
          });
        });
      }
    }

    // ── Handle messages from content scripts ──────────────────────────────
    chrome.runtime.onMessage.addListener((msg, sender, sendResponse) => {
      if (msg.type === 'analyze_claims') {
        const { claims, platform } = msg;
        const sent = sendToBackend({
          client_id: CLIENT_ID,
          claims,
          platform,
          timestamp: Date.now() / 1000,
        });
        sendResponse({ queued: sent });
        return true;
      }

      if (msg.type === 'get_ws_status') {
        sendResponse({ status: isConnected ? 'connected' : 'offline' });
        return true;
      }

      if (msg.type === 'get_client_id') {
        sendResponse({ clientId: CLIENT_ID });
        return true;
      }

      if (msg.type === 'reconnect') {
        if (!isConnected) connect();
        sendResponse({ ok: true });
        return true;
      }
    });

    // ── Init ──────────────────────────────────────────────────────────────
    connect();

    // Keep alive: ping every 20s to prevent Manifest V3 service worker suspension
    setInterval(() => {
      if (ws && ws.readyState === WebSocket.OPEN) {
        try {
          ws.send(JSON.stringify({ type: 'ping', client_id: CLIENT_ID }));
        } catch (_) {
          connect();
        }
      } else if (!ws || ws.readyState === WebSocket.CLOSED) {
        connect();
      }
    }, 20_000);
  },
});

function getOrCreateClientId(): string {
  // In background context, use a stable ID from chrome.storage.local
  // Synchronous approximation; actual persistence handled by store
  const stored = sessionStorage.getItem('fact_engine_client_id');
  if (stored) return stored;
  const arr = new Uint8Array(8);
  crypto.getRandomValues(arr);
  const id = 'ext-' + Array.from(arr).map(b => b.toString(16).padStart(2, '0')).join('');
  sessionStorage.setItem('fact_engine_client_id', id);
  return id;
}