LexAI/entrypoints/background.ts

import { defineBackground } from 'wxt/utils/define-background';
import type { AnalyzePayload, LexAIConfig, LexAIResponse } from '@lib/types';
import { ACTIONS, ACTION_LABELS, CONTEXT_MENU_STYLES } from '@lib/actions';
import { decryptApiKey } from '@lib/crypto';
import { callProvider, getSystemPrompt, listModels } from '@lib/providers';

// Resolve the usable API key from stored config: prefer the encrypted path,
// fall back to plaintext for backward compat. Returns null if none is set.
async function resolveApiKey(config: LexAIConfig): Promise<string | null> {
  let apiKey = config.apiKey;
  if (config.apiKeyEnc && config.encKey) {
    const decrypted = decryptApiKey(config.encKey, config.apiKeyEnc);
    if (decrypted) apiKey = decrypted;
  }
  if (!apiKey || apiKey.trim() === '') return null;
  return apiKey.trim();
}

// ─── Main handler ─────────────────────────────────────────────────────────────

async function handleAnalyzeText(payload: AnalyzePayload): Promise<LexAIResponse> {
  const stored = await chrome.storage.local.get(['provider', 'apiKey', 'apiKeyEnc', 'encKey', 'model']);
  const config = stored as LexAIConfig;

  const apiKey = await resolveApiKey(config);
  if (!apiKey) {
    return { error: 'No API key configured. Please open LexAI settings (click the extension icon).' };
  }

  const resolvedConfig: LexAIConfig = { ...config, apiKey };
  const systemPrompt = getSystemPrompt(payload.action, payload.style);
  return callProvider(resolvedConfig, payload.text, systemPrompt);
}

// ─── Background entry ─────────────────────────────────────────────────────────

export default defineBackground(() => {
  console.log('LexAI background service worker started');

  // ─── Context menus ───────────────────────────────────────────────────────
  chrome.runtime.onInstalled.addListener(() => {
    chrome.contextMenus.removeAll(() => {
      ACTIONS.forEach(action => {
        chrome.contextMenus.create({
          id: `lexai-${action}`,
          title: `⚡ LexAI: ${ACTION_LABELS[action]}`,
          contexts: ['selection'],
        });
        CONTEXT_MENU_STYLES.forEach(style => {
          chrome.contextMenus.create({
            id: `lexai-${action}-${style.toLowerCase()}`,
            parentId: `lexai-${action}`,
            title: style,
            contexts: ['selection'],
          });
        });
      });
    });
  });

  chrome.contextMenus.onClicked.addListener((info, tab) => {
    if (!info.selectionText || !tab?.id) return;
    // Parse action and style from menuItemId e.g. "lexai-fix-formal"
    const parts = info.menuItemId.toString().replace('lexai-', '').split('-');
    const action = parts[0];
    const style = parts[1] ? parts[1].charAt(0).toUpperCase() + parts[1].slice(1) : 'Default';
    chrome.tabs.sendMessage(tab.id, {
      type: 'lexai-context-menu',
      action,
      text: info.selectionText,
      style,
    });
  });

  // ─── Message handler ─────────────────────────────────────────────────────
  chrome.runtime.onMessage.addListener((message, _sender, sendResponse) => {
    if (message.type === 'ANALYZE_TEXT') {
      // Support both { payload: { text, action, style } } (content.ts) and
      // { text, action, style } (popup) formats
      const payload: AnalyzePayload = message.payload ?? {
        text: message.text as string,
        action: message.action as string,
        style: message.style as string | undefined,
      };
      handleAnalyzeText(payload)
        .then(sendResponse)
        .catch((err) => sendResponse({ error: String(err) }));
      return true; // Keep channel open for async response
    }

    if (message.type === 'COPY_AS') {
      const { text, format } = message as { text: string; format: string };
      chrome.storage.local.get(['provider', 'apiKey', 'apiKeyEnc', 'encKey', 'model'])
        .then(async (stored) => {
          const config = stored as LexAIConfig;
          const apiKey = await resolveApiKey(config);
          if (!apiKey) {
            sendResponse({ error: 'No API key configured. Please open LexAI settings.' });
            return;
          }
          const resolvedConfig: LexAIConfig = { ...config, apiKey };
          const systemPrompt = `Reformat the following text as ${format}. Return only the reformatted result, no explanation.`;
          return callProvider(resolvedConfig, text, systemPrompt);
        })
        .then((res) => { if (res) sendResponse(res); })
        .catch((err) => sendResponse({ error: String(err) }));
      return true;
    }

    if (message.type === 'LIST_MODELS') {
      const provider = (message.provider as string) || 'openai';
      // Prefer an inline key (freshly typed, not yet saved); else use the stored key.
      const inlineKey = (message.apiKey as string | undefined)?.trim() || undefined;
      (async () => {
        let apiKey = inlineKey;
        if (!apiKey) {
          const stored = await chrome.storage.local.get(['apiKey', 'apiKeyEnc', 'encKey']);
          apiKey = (await resolveApiKey(stored as LexAIConfig)) ?? undefined;
        }
        if (!apiKey && provider !== 'openrouter') {
          return { error: 'No API key found. Enter your API key above, then click Load models.' };
        }
        return listModels(provider, apiKey);
      })()
        .then(sendResponse)
        .catch((err) => sendResponse({ error: String(err) }));
      return true;
    }
  });
});