background.js

/**
 * RPC Extension Background Script
 * Main entry point that coordinates function execution via broadcast messages
 */

// Import library modules (loaded via importScripts in service worker)
importScripts(
  'lib/config.js',
  'lib/functionRegistry.js',
  'lib/centrifugoClient.js'
);

// Global state
let extensionInitialized = false;
let initializationPromise = null;  // Prevent concurrent initialization
let centrifugoClient = null;
const pendingRequests = new Map(); // Track pending requests with timeouts

/**
 * Check RPC service connection status
 */
async function checkRpcServiceConnection() {
  if (!config.rpcServiceUrl) {
    console.warn('[RPC Extension] RPC_SERVICE_URL not configured');
    return { connected: false, error: 'URL not configured' };
  }

  try {
    const response = await fetch(`${config.rpcServiceUrl}/health`, {
      method: 'GET',
      headers: { 'Accept': 'application/json' }
    });

    if (response.ok) {
      const data = await response.json();
      console.log('[RPC Extension] ✓ RPC Service connected:', config.rpcServiceUrl);
      console.log('[RPC Extension]   Status:', data.status);
      console.log('[RPC Extension]   Version:', data.version);
      return { connected: true, status: data };
    } else {
      console.error('[RPC Extension] ✗ RPC Service returned:', response.status);
      return { connected: false, error: `HTTP ${response.status}` };
    }
  } catch (error) {
    console.error('[RPC Extension] ✗ RPC Service unreachable:', error.message);
    return { connected: false, error: error.message };
  }
}

/**
 * Initialize Centrifugo WebSocket connection for low-latency RPC
 */
async function initializeCentrifugo() {
  centrifugoClient = new CentrifugoClient({
    wsUrl: config.centrifugoWsUrl,
    token: config.centrifugoToken,
    sessionId: config.sessionId,

    onConnect: () => {},

    onDisconnect: (info) => {
      console.warn('[RPC Extension] Disconnected:', info.reason || 'unknown');
      // Notify content scripts to stop the automation glow effect
      browser.tabs.sendMessage(null, { type: 'stop_automation_effects' }).catch(() => {});
    },

    onError: (error) => {
      console.error('[RPC Extension] Connection error:', error);
    },

    onMessage: async (data) => {
      if (data && data.type === 'rpc_call') {
        await executeRpcCall(data);
      }
    }
  });

  await centrifugoClient.connect();
}

/**
 * Execute an RPC call (from Centrifugo or broadcast)
 */
async function executeRpcCall(message) {
  const { id, method, parameters, timeout = 30000 } = message;

  if (!id || !method) {
    console.error('[RPC Extension] Invalid RPC message:', message);
    return;
  }

  // Check if we handle this function
  if (!functionRegistry.has(method)) {
    // Forward to content scripts for other extensions to handle
    browser.tabs.sendMessage(null, {
      type: 'execute_effect',
      method: method,
      params: parameters || {},
      id: id,
      rpc_config: {
        rpcServiceUrl: config.rpcServiceUrl,
        sessionId: config.sessionId,
        rpcApiKey: config.rpcApiKey
      }
    });
    return;
  }

  // Set up timeout
  const timeoutId = setTimeout(async () => {
    console.error(`[RPC Extension] Request ${id} timed out after ${timeout}ms`);
    pendingRequests.delete(id);
    // Post timeout via Centrifugo (or fallback to HTTP)
    await postResultViaCentrifugo(id, 'timeout', null, `Request timed out after ${timeout}ms`);
  }, timeout);

  pendingRequests.set(id, timeoutId);

  try {
    // Execute the function with context
    const result = await functionRegistry.execute(method, parameters, { id });

    // Clear timeout
    clearTimeout(timeoutId);
    pendingRequests.delete(id);

    // Check if this was delegated to content script
    if (result && result.delegated) {
      return;
    }

    // Post success result via Centrifugo
    await postResultViaCentrifugo(id, 'success', result, null);

  } catch (error) {
    // Clear timeout
    clearTimeout(timeoutId);
    pendingRequests.delete(id);

    // Post error result via Centrifugo
    console.error(`[RPC Extension] Execution error for ${id}:`, error);
    await postResultViaCentrifugo(id, 'error', null, error.message || String(error));
  }
}

// Chunk size for large payloads (1MB - requires Centrifugo websocket.message_size_limit: 1048576)
const CHUNK_SIZE = 1024 * 1024;

/**
 * Post RPC result via Centrifugo, chunking if necessary
 */
async function postResultViaCentrifugo(uniqueId, status, result, error) {
  if (!centrifugoClient || !centrifugoClient.isConnected()) {
    console.error(`[RPC Extension] Cannot post result - Centrifugo not connected`);
    return;
  }

  const resultPayload = {
    id: uniqueId,
    status: status,
    result: result,
    error: error
  };

  const payloadStr = JSON.stringify(resultPayload);
  const payloadSize = payloadStr.length;

  // If payload fits in one message, send directly
  if (payloadSize <= CHUNK_SIZE) {
    const success = await centrifugoClient.publishResult(resultPayload);
    if (!success) {
      console.error(`[RPC Extension] Failed to post result for ${uniqueId}`);
    }
    return;
  }

  // Chunk the payload for large results
  const totalChunks = Math.ceil(payloadSize / CHUNK_SIZE);

  for (let i = 0; i < totalChunks; i++) {
    const start = i * CHUNK_SIZE;
    const end = Math.min(start + CHUNK_SIZE, payloadSize);
    const chunkData = payloadStr.substring(start, end);

    const chunkPayload = {
      id: uniqueId,
      chunked: true,
      chunk_index: i,
      total_chunks: totalChunks,
      data: chunkData
    };

    const success = await centrifugoClient.publishResult(chunkPayload);
    if (!success) {
      console.error(`[RPC Extension] Failed to send chunk ${i + 1}/${totalChunks} for ${uniqueId}`);
      return;
    }
  }
}

/**
 * Initialize the extension
 */
async function initializeExtension() {
  // Return existing promise if initialization is in progress
  if (initializationPromise) {
    return initializationPromise;
  }
  if (extensionInitialized) {
    return;
  }
  // Start initialization and store the promise
  initializationPromise = _doInitialize();
  return initializationPromise;
}

async function _doInitialize() {
  try {
    console.log('[RPC Extension] Initializing...');

    // Initialize configuration (will throw if auth fails)
    await config.init();

    // Check RPC service connection (non-blocking)
    checkRpcServiceConnection().catch(() => {});

    // Initialize Centrifugo client (required for communication)
    await initializeCentrifugo();

    // Register automation proxies (delegate to content script)
    registerAutomationProxies();

    extensionInitialized = true;

    // Register Session API methods
    registerSessionApiMethods();

    // Expose registry to global scope for manual registration
    // Use webfuse.rpc namespace, preserving any existing webfuse global
    globalThis.webfuse = globalThis.webfuse || {};
    globalThis.webfuse.rpc = globalThis.webfuse.rpc || {};
    Object.assign(globalThis.webfuse.rpc, {
      register: (name, handler, schema) => {
        functionRegistry.register(name, handler, schema);
      },
      unregister: (name) => functionRegistry.unregister(name),
      list: () => functionRegistry.list(),
      getSchema: (name) => functionRegistry.getSchema(name),
      getAllSchemas: () => functionRegistry.getAllSchemas()
    });

  } catch (error) {
    console.error('[RPC Extension] Initialization failed:', error);
    throw error;
  }
}

/**
 * Register Session API methods
 */
function registerSessionApiMethods() {
  // Listen for session events via broadcast listener
  browser.webfuseSession.on.addListener((message, sender) => {
    // Check if this is a session_ended event
    if (message?.event_type === 'session_ended') {
      browser.tabs.sendMessage(null, { type: 'stop_automation_effects' }).catch(() => {});
    }
  });

  // Register relocate method for navigation
  functionRegistry.register('relocate', async (params) => {
    const { url, newTab = false } = params;
    browser.webfuseSession.relocate(url, newTab);
    return { success: true, url };
  }, {
    params: [
      { name: 'url', type: 'string', description: 'The URL to navigate to (e.g., "https://example.com")', required: true },
      { name: 'newTab', type: 'boolean', description: 'If true, open URL in a new tab instead of current tab', required: false }
    ],
    returns: 'object',
    description: 'Navigate to a URL. Use this to go to a new page. Set newTab=true to open in a new tab while keeping the current page.'
  });

  // Register openTab method
  functionRegistry.register('open_tab', async (params) => {
    const { url } = params;
    browser.webfuseSession.openTab(url);
    return { success: true, url };
  }, {
    params: [
      { name: 'url', type: 'string', description: 'The URL to open in the new tab', required: true }
    ],
    returns: 'object',
    description: 'Open a new browser tab with the specified URL. The current tab remains unchanged.'
  });
}

/**
 * Invalidate all pending requests (called on navigation)
 */
function invalidatePendingRequests() {
  // Cancel all pending RPC request timeouts
  for (const [uniqueId, timeoutId] of pendingRequests.entries()) {
    clearTimeout(timeoutId);
    pendingRequests.delete(uniqueId);
    // Post error for each invalidated request via Centrifugo
    postResultViaCentrifugo(uniqueId, 'error', null, 'Request cancelled due to page navigation').catch(err => {
      console.error(`[RPC Extension] Failed to post error for invalidated request ${uniqueId}:`, err);
    });
  }
}

/**
 * Register proxy functions that delegate to content script for automation
 * Parameter schemas based on Webfuse Automation API documentation
 */
function registerAutomationProxies() {
  // Define detailed schemas for each automation method
  // Based on https://docs.webfuse.io/extensions/automation/automation-api
  const automationSchemas = [
    {
      name: 'mouse_move',
      description: 'Move the virtual mouse pointer to a target element or coordinates. Use CSS selectors like "#id" or ".class" to target elements.',
      params: [
        { name: 'target', type: 'string', description: 'CSS selector (e.g., "#button", ".input-field") or coordinates [x, y]', required: true },
        { name: 'persistent', type: 'boolean', description: 'Keep mouse pointer visible on screen (hides after some time by default)', required: false }
      ]
    },
    {
      name: 'scroll',
      description: 'Scroll a scrollable element by a specified amount in a direction. Use to navigate long pages or scroll containers.',
      params: [
        { name: 'target', type: 'string', description: 'CSS selector of the scrollable element (e.g., "#content", "body")', required: true },
        { name: 'direction', type: 'string', description: 'Scroll direction: "vertical" or "horizontal"', required: true },
        { name: 'amount', type: 'number', description: 'Number of pixels to scroll (positive = down/right, negative = up/left)', required: true }
      ]
    },
    {
      name: 'left_click',
      description: 'Perform a left mouse button click on an element. Use CSS selectors to target buttons, links, or any clickable elements.',
      params: [
        { name: 'target', type: 'string', description: 'CSS selector of the element to click (e.g., "#submit-btn", "button.primary")', required: true },
        { name: 'moveMouse', type: 'boolean', description: 'Move virtual mouse pointer to target before clicking', required: false }
      ]
    },
    {
      name: 'middle_click',
      description: 'Perform a middle mouse button (wheel) click. Often used to open links in new tabs.',
      params: [
        { name: 'target', type: 'string', description: 'CSS selector of the element to click', required: true },
        { name: 'moveMouse', type: 'boolean', description: 'Move virtual mouse pointer to target before clicking', required: false }
      ]
    },
    {
      name: 'right_click',
      description: 'Perform a right mouse button click to open context menus.',
      params: [
        { name: 'target', type: 'string', description: 'CSS selector of the element to right-click', required: true },
        { name: 'moveMouse', type: 'boolean', description: 'Move virtual mouse pointer to target before clicking', required: false }
      ]
    },
    {
      name: 'type',
      description: 'Type text into an input field or text area. Simulates natural human typing. Use this to fill forms, search boxes, or any text input.',
      params: [
        { name: 'target', type: 'string', description: 'CSS selector of the input element (e.g., "input[name=email]", "#search-box", "textarea.comment")', required: true },
        { name: 'text', type: 'string', description: 'The text content to type into the input field', required: true },
        { name: 'moveMouse', type: 'boolean', description: 'Move virtual mouse pointer to target before typing', required: false },
        { name: 'overwrite', type: 'boolean', description: 'Clear existing content before typing (select all + replace)', required: false },
        { name: 'timePerChar', type: 'number', description: 'Average milliseconds between keystrokes (default: 100)', required: false }
      ]
    },
    {
      name: 'key_press',
      description: 'Press a keyboard key, optionally with modifier keys (Ctrl, Alt, Shift, Meta). Use for keyboard shortcuts or special keys like Enter, Escape, Tab.',
      params: [
        { name: 'target', type: 'string', description: 'CSS selector of the element to receive the key press', required: true },
        { name: 'key', type: 'string', description: 'Key to press (e.g., "Enter", "Escape", "Tab", "a", "A", "ArrowDown")', required: true },
        { name: 'options', type: 'object', description: 'Modifier keys: {altKey: bool, ctrlKey: bool, metaKey: bool, shiftKey: bool}', required: false }
      ]
    },
    {
      name: 'take_dom_snapshot',
      description: 'Capture the current page DOM structure for analysis. Returns a text representation of the page that can be used to understand page structure and find CSS selectors for automation. Use this before interacting with elements to understand the page layout.',
      params: [
        { name: 'options', type: 'object', description: 'Snapshot options: {rootSelector: string, crossframe: bool, modifier: "downsample"|"full"|{name: string, params: object}}', required: false }
      ]
    },
    {
      name: 'take_gui_snapshot',
      description: 'Take a screenshot of the current page. Returns an image that shows the visual state of the page. Use this to see what the user sees and verify visual elements.',
      params: [
        { name: 'options', type: 'object', description: 'Screenshot options (currently unused)', required: false }
      ]
    }
  ];

  // Methods that may trigger navigation and won't be able to post result
  const navigationTriggeringMethods = ['left_click', 'middle_click', 'right_click', 'key_press'];

  // Register each automation method
  automationSchemas.forEach(schema => {
    const handler = async (params, context) => {
      const mayTriggerNavigation = navigationTriggeringMethods.includes(schema.name);

      // Send message to content script
      browser.tabs.sendMessage(null, {
        type: 'execute_automation',
        method: schema.name,
        params,
        id: context?.id,
        rpc_config: {
          rpcServiceUrl: config.rpcServiceUrl,
          sessionId: config.sessionId,
          rpcApiKey: config.rpcApiKey
        },
        // Tell content script NOT to post result if this may trigger navigation
        skipResultPost: mayTriggerNavigation
      });

      // For navigation-triggering methods, post success immediately from background
      // because the content script may be destroyed before it can post
      if (mayTriggerNavigation) {
        // Small delay to allow the click to be dispatched
        await new Promise(resolve => setTimeout(resolve, 100));
        return { success: true, action: schema.name };
      }

      // Return delegated for non-navigation methods - content script will post result
      return { delegated: true };
    };

    functionRegistry.register(schema.name, handler, {
      params: schema.params,
      returns: 'any',
      description: schema.description
    });
  });

  // Register wait() locally (not delegated to content script)
  functionRegistry.register('wait', async (params) => {
    const { ms } = params;
    if (typeof ms !== 'number' || ms < 0) {
      throw new Error('wait() requires a non-negative number of milliseconds');
    }
    return new Promise(resolve => setTimeout(resolve, ms));
  }, {
    params: [
      { name: 'ms', type: 'number', description: 'Number of milliseconds to wait', required: true }
    ],
    returns: 'void',
    description: 'Wait for a specified amount of time. Use this to pause between actions or wait for page elements to load.'
  });
}


/**
 * Handle function discovery requests
 */
async function handleDiscoveryRequest(message, sender) {
  if (!message || message.type !== 'rpc_discovery') {
    return;
  }

  const { id } = message;

  try {
    const functions = functionRegistry.getAllSchemas();

    // Post discovery result via Centrifugo
    await postResultViaCentrifugo(id, 'success', { functions }, null);

  } catch (error) {
    console.error('[RPC Extension] Discovery error:', error);
    await postResultViaCentrifugo(id, 'error', null, error.message || String(error));
  }
}

/**
 * Handle messages from content script (automation results and RPC result forwarding)
 */
browser.runtime.onMessage.addListener((message, sender) => {
  // Handle automation execution responses from content script
  if (message.type === 'automation_result') {
    return Promise.resolve(message);
  }

  // Handle RPC result forwarding from content script
  // This is used when other extensions post results via DOM events
  if (message.type === 'post_rpc_result') {
    const { id, status, result, error } = message;

    // Post result via Centrifugo WebSocket
    (async () => {
      await postResultViaCentrifugo(id, status, result, error);
    })();

    return false;
  }

  return false;
});

/**
 * Validate message structure for security
 */
function isValidRpcMessage(message) {
  // Must be an object
  if (!message || typeof message !== 'object') {
    console.warn('[RPC Extension] Rejected: message is not an object');
    return false;
  }

  // Must have a valid type
  if (message.type !== 'rpc_call' && message.type !== 'rpc_discovery') {
    // Not an RPC message - ignore silently (may be intended for other handlers)
    return false;
  }

  // Must have id for RPC calls
  if (message.type === 'rpc_call') {
    if (!message.id || typeof message.id !== 'string') {
      console.warn('[RPC Extension] Rejected: rpc_call missing valid id');
      return false;
    }
    if (!message.method || typeof message.method !== 'string') {
      console.warn('[RPC Extension] Rejected: rpc_call missing valid method');
      return false;
    }
    // Validate method name format (alphanumeric + underscores)
    if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(message.method)) {
      console.warn('[RPC Extension] Rejected: invalid method name format');
      return false;
    }
    // Validate parameters if present
    if (message.parameters !== undefined && typeof message.parameters !== 'object') {
      console.warn('[RPC Extension] Rejected: parameters must be an object');
      return false;
    }
  }

  if (message.type === 'rpc_discovery') {
    if (!message.id || typeof message.id !== 'string') {
      console.warn('[RPC Extension] Rejected: rpc_discovery missing valid id');
      return false;
    }
  }

  return true;
}

/**
 * Listen for broadcast messages
 */
browser.webfuseSession.on.addListener(async (message, sender) => {
  if (!extensionInitialized) {
    await initializeExtension();
  }

  // Unwrap the message - Webfuse wraps it in a "message" property
  const unwrappedMessage = message?.message || message;

  // Validate message structure
  if (!isValidRpcMessage(unwrappedMessage)) {
    return;
  }

  // Handle different message types
  if (unwrappedMessage.type === 'rpc_call') {
    await executeRpcCall(unwrappedMessage);
  } else if (unwrappedMessage.type === 'rpc_discovery') {
    await handleDiscoveryRequest(unwrappedMessage, sender);
  }
});

/**
 * Initialize on extension load
 */
(async () => {
  try {
    await initializeExtension();
    console.log('[RPC Extension] Ready to receive RPC calls');
  } catch (error) {
    console.error('[RPC Extension] Startup failed:', error);
  }
})();