LexAI/packages/chrome/entrypoints/background.ts

import { defineBackground } from 'wxt/utils/define-background';
import { CONFIG_STORAGE_KEYS } from '@lib/types';
import type { AnalyzePayload, LexAIConfig, LexAIResponse } from '@lib/types';
import { CONTEXT_MENU_ENTRIES, findContextMenuEntry, resolvePromptPersona, resolvePromptPattern } from '@lib/actions';
import { decryptApiKey, migratePlaintextApiKey } from '@lib/crypto';
import { callProvider, defaultMaxTokens, getSystemPrompt, listModels, providerLabel } 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();
}

// Options can save a provider change without touching the stored key, which
// leaves e.g. an OpenAI key attached to Groq — every call then fails with the
// provider's own "Invalid API Key". Detect it here and say what to do instead.
// `keyProvider` is absent for keys saved by older builds: unknown, so allow it.
function keyProviderMismatch(config: LexAIConfig, provider: string): string | null {
  if (!config.keyProvider || config.keyProvider === provider) return null;
  return `Your saved API key was entered for ${providerLabel(config.keyProvider)}, but the selected provider is ${providerLabel(provider)}. Open LexAI Settings and enter a ${providerLabel(provider)} API key.`;
}

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

async function handleAnalyzeText(payload: AnalyzePayload): Promise<LexAIResponse> {
  const stored = await chrome.storage.local.get([...CONFIG_STORAGE_KEYS]);
  const config = stored as LexAIConfig;

  // Prompt requests from the toolbar/context menu carry no explicit params —
  // apply the Prompt Builder settings saved from the popup so all entry
  // points behave the same. The popup still overrides by sending its own.
  if (payload.action === 'prompt' && !payload.promptParams) {
    const saved = (await chrome.storage.local.get([
      'promptPattern', 'promptStyle', 'promptPersona', 'customPersona', 'promptFormat', 'promptModel',
    ])) as Record<string, string | undefined>;
    payload = {
      ...payload,
      promptParams: {
        pattern: resolvePromptPattern(saved.promptPattern, saved.promptStyle),
        persona: resolvePromptPersona(saved.promptPersona, saved.customPersona),
        format: saved.promptFormat,
      },
      model: payload.model ?? (saved.promptModel || undefined),
    };
  }

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

  const mismatch = keyProviderMismatch(config, config.provider || 'openai');
  if (mismatch) return { error: mismatch };

  const resolvedConfig: LexAIConfig = {
    ...config,
    apiKey,
    ...(payload.model ? { model: payload.model } : {}),
  };
  const systemPrompt = getSystemPrompt(payload.action, payload.style, payload.promptParams);
  // The Prompt Builder's engineered prompts (react/pev/gauntlet skeletons
  // especially) run well past defaultMaxTokens' input-scaled floor for a
  // short input idea — give the 'prompt' action a higher floor.
  const callOpts = payload.action === 'prompt'
    ? { maxTokens: Math.max(2048, defaultMaxTokens(payload.text)) }
    : undefined;
  return callProvider(resolvedConfig, payload.text, systemPrompt, callOpts);
}

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

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

  // Encrypt any legacy plaintext apiKey left by older builds (no-op otherwise).
  // The read path keeps its plaintext fallback, so a failed migration is safe.
  migratePlaintextApiKey().catch(() => {});

  // ─── Context menus ───────────────────────────────────────────────────────
  chrome.runtime.onInstalled.addListener(() => {
    chrome.contextMenus.removeAll(() => {
      CONTEXT_MENU_ENTRIES.forEach((entry) => {
        chrome.contextMenus.create({
          id: entry.id,
          parentId: entry.parentId,
          title: entry.title,
          contexts: ['selection'],
        });
      });
    });
  });

  chrome.contextMenus.onClicked.addListener((info, tab) => {
    if (!info.selectionText || !tab?.id) return;
    const entry = findContextMenuEntry(info.menuItemId.toString());
    if (!entry) return;
    chrome.tabs.sendMessage(tab.id, {
      type: 'lexai-context-menu',
      action: entry.action,
      text: info.selectionText,
      style: entry.style,
    });
  });

  // ─── Message handler ─────────────────────────────────────────────────────
  chrome.runtime.onMessage.addListener((message, sender, sendResponse) => {
    // Only our own contexts (content scripts, popup, options) may drive the
    // key-bearing call path — ignore anything from another extension.
    if (sender.id !== chrome.runtime.id) return;

    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,
        promptParams: message.promptParams,
        model: message.model 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([...CONFIG_STORAGE_KEYS])
        .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 mismatch = keyProviderMismatch(config, config.provider || 'openai');
          if (mismatch) {
            sendResponse({ error: mismatch });
            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') {
      // Prefer an inline key (freshly typed, not yet saved); else use the stored key.
      const inlineKey = (message.apiKey as string | undefined)?.trim() || undefined;
      (async () => {
        // Callers that don't know the provider (content script) omit it —
        // fall back to the configured one.
        let provider = message.provider as string | undefined;
        if (!provider) {
          const stored = await chrome.storage.local.get('provider');
          provider = (stored.provider as string) || 'openai';
        }
        let apiKey = inlineKey;
        if (!apiKey) {
          // Stored-key path: refuse to send a key that belongs to a different
          // provider — it would come back as that provider's "invalid key".
          const stored = await chrome.storage.local.get(['apiKey', 'apiKeyEnc', 'encKey', 'keyProvider']);
          const mismatch = keyProviderMismatch(stored as LexAIConfig, provider);
          if (mismatch) return { error: mismatch, keyRejected: true };
          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;
    }
  });
});