LexAI/entrypoints/background.ts

import { defineBackground } from 'wxt/utils/define-background';
import nacl from 'tweetnacl';

// ─── Types ────────────────────────────────────────────────────────────────────

interface AnalyzePayload {
  text: string;
  action: string;
}

interface LexAIConfig {
  provider?: string;
  apiKey?: string;
  apiKeyEnc?: string;
  encKey?: string;
  model?: string;
}

interface LexAIResponse {
  result?: string;
  error?: string;
}

// ─── System prompts ───────────────────────────────────────────────────────────

function getSystemPrompt(action: string): string {
  const prompts: Record<string, string> = {
    grammar:
      'You are a professional grammar editor. Fix all grammar, spelling, and punctuation errors in the provided text. ' +
      'Preserve the original meaning and tone as closely as possible. ' +
      'Return ONLY the corrected text — no explanations, no preamble.',
    rephrase:
      'You are a skilled writing assistant. Rephrase the provided text to make it clearer, more engaging, and more professional. ' +
      'Keep the same meaning and approximate length. ' +
      'Return ONLY the rephrased text — no explanations.',
    shorten:
      'You are a concise editor. Shorten the provided text by at least 30% while preserving the core message. ' +
      'Remove filler words, redundant phrases, and unnecessary detail. ' +
      'Return ONLY the shortened text.',
    expand:
      'You are an experienced writer. Expand the provided text with more detail, context, and supporting points. ' +
      'Make it richer and more informative while staying on topic. ' +
      'Return ONLY the expanded text.',
    tone:
      'You are a writing coach. Analyze the tone of the provided text. ' +
      'Describe the tone characteristics (e.g. formal, casual, aggressive, passive, confident, etc.) ' +
      'and note any issues like passive voice, wordiness, or emotional bias. ' +
      'Return a brief, clear analysis — no rewriting, no extra commentary.',
  };
  return prompts[action] ?? prompts.grammar;
}

// ─── Encryption helpers ───────────────────────────────────────────────────────

async function decryptApiKey(encKeyB64: string, apiKeyEncB64: string): Promise<string | null> {
  const key      = Uint8Array.from(atob(encKeyB64),   c => c.charCodeAt(0));
  const combined = Uint8Array.from(atob(apiKeyEncB64), c => c.charCodeAt(0));
  const nonce    = combined.slice(0, 24);
  const cipher   = combined.slice(24);
  const decrypted = nacl.secretbox.open(cipher, nonce, key);
  if (!decrypted) return null;
  return new TextDecoder().decode(decrypted);
}

// ─── Provider implementations ─────────────────────────────────────────────────

async function callOpenAI(payload: AnalyzePayload, config: LexAIConfig): Promise<LexAIResponse> {
  const model = config.model || 'gpt-4o-mini';
  let res: Response;

  try {
    res = await fetch('https://api.openai.com/v1/chat/completions', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        Authorization: `Bearer ${config.apiKey}`,
      },
      body: JSON.stringify({
        model,
        messages: [
          { role: 'system', content: getSystemPrompt(payload.action) },
          { role: 'user', content: payload.text },
        ],
        max_tokens: 1024,
        temperature: 0.7,
      }),
    });
  } catch (err) {
    return { error: `Network error reaching OpenAI: ${String(err)}` };
  }

  const data = await res.json();

  if (!res.ok) {
    const msg = data?.error?.message ?? `HTTP ${res.status}`;
    return { error: `OpenAI error: ${msg}` };
  }

  const result = data?.choices?.[0]?.message?.content as string | undefined;
  if (!result) return { error: 'OpenAI returned an empty response.' };
  return { result: result.trim() };
}

async function callAnthropic(payload: AnalyzePayload, config: LexAIConfig): Promise<LexAIResponse> {
  const model = config.model || 'claude-3-5-haiku-20241022';
  let res: Response;

  try {
    res = await fetch('https://api.anthropic.com/v1/messages', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'x-api-key': config.apiKey!,
        'anthropic-version': '2023-06-01',
      },
      body: JSON.stringify({
        model,
        max_tokens: 1024,
        system: getSystemPrompt(payload.action),
        messages: [{ role: 'user', content: payload.text }],
      }),
    });
  } catch (err) {
    return { error: `Network error reaching Anthropic: ${String(err)}` };
  }

  const data = await res.json();

  if (!res.ok) {
    const msg = data?.error?.message ?? `HTTP ${res.status}`;
    return { error: `Anthropic error: ${msg}` };
  }

  const result = data?.content?.[0]?.text as string | undefined;
  if (!result) return { error: 'Anthropic returned an empty response.' };
  return { result: result.trim() };
}

async function callGroq(payload: AnalyzePayload, config: LexAIConfig): Promise<LexAIResponse> {
  const model = config.model || 'llama-3.3-70b-versatile';
  let res: Response;

  try {
    res = await fetch('https://api.groq.com/openai/v1/chat/completions', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        Authorization: `Bearer ${config.apiKey}`,
      },
      body: JSON.stringify({
        model,
        messages: [
          { role: 'system', content: getSystemPrompt(payload.action) },
          { role: 'user', content: payload.text },
        ],
        max_tokens: 1024,
        temperature: 0.7,
      }),
    });
  } catch (err) {
    return { error: `Network error reaching Groq: ${String(err)}` };
  }

  const data = await res.json();

  if (!res.ok) {
    const msg = data?.error?.message ?? `HTTP ${res.status}`;
    return { error: `Groq error: ${msg}` };
  }

  const result = data?.choices?.[0]?.message?.content as string | undefined;
  if (!result) return { error: 'Groq returned an empty response.' };
  return { result: result.trim() };
}

async function callOpenRouter(payload: AnalyzePayload, config: LexAIConfig): Promise<LexAIResponse> {
  const model = config.model || 'openai/gpt-4o-mini';
  let res: Response;

  try {
    res = await fetch('https://openrouter.ai/api/v1/chat/completions', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        Authorization: `Bearer ${config.apiKey}`,
        'HTTP-Referer': 'https://lexai.dev',
        'X-Title': 'LexAI',
      },
      body: JSON.stringify({
        model,
        messages: [
          { role: 'system', content: getSystemPrompt(payload.action) },
          { role: 'user', content: payload.text },
        ],
        max_tokens: 1024,
      }),
    });
  } catch (err) {
    return { error: `Network error reaching OpenRouter: ${String(err)}` };
  }

  const data = await res.json();

  if (!res.ok) {
    const msg = data?.error?.message ?? `HTTP ${res.status}`;
    return { error: `OpenRouter error: ${msg}` };
  }

  const result = data?.choices?.[0]?.message?.content as string | undefined;
  if (!result) return { error: 'OpenRouter returned an empty response.' };
  return { result: result.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;

  // Resolve API key — prefer encrypted path, fall back to plaintext for backward compat
  let apiKey = config.apiKey;
  if (config.apiKeyEnc && config.encKey) {
    const decrypted = await decryptApiKey(config.encKey, config.apiKeyEnc);
    if (decrypted) {
      apiKey = decrypted;
    }
  }

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

  const resolvedConfig: LexAIConfig = { ...config, apiKey: apiKey.trim() };
  const provider = config.provider || 'openai';

  switch (provider) {
    case 'openai':
      return callOpenAI(payload, resolvedConfig);
    case 'anthropic':
      return callAnthropic(payload, resolvedConfig);
    case 'groq':
      return callGroq(payload, resolvedConfig);
    case 'openrouter':
      return callOpenRouter(payload, resolvedConfig);
    default:
      return { error: `Unknown provider: "${provider}". Please check LexAI settings.` };
  }
}

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

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

  // ─── Context menus ───────────────────────────────────────────────────────
  chrome.runtime.onInstalled.addListener(() => {
    chrome.contextMenus.create({
      id: 'lexai-grammar',
      title: '⚡ LexAI: Fix Grammar',
      contexts: ['selection'],
    });
    chrome.contextMenus.create({
      id: 'lexai-rephrase',
      title: '⚡ LexAI: Rephrase',
      contexts: ['selection'],
    });
    chrome.contextMenus.create({
      id: 'lexai-shorten',
      title: '⚡ LexAI: Shorten',
      contexts: ['selection'],
    });
    chrome.contextMenus.create({
      id: 'lexai-expand',
      title: '⚡ LexAI: Expand',
      contexts: ['selection'],
    });
    chrome.contextMenus.create({
      id: 'lexai-tone',
      title: '⚡ LexAI: Analyze Tone',
      contexts: ['selection'],
    });
  });

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

  // ─── Message handler ─────────────────────────────────────────────────────
  chrome.runtime.onMessage.addListener((message, _sender, sendResponse) => {
    if (message.type === 'ANALYZE_TEXT') {
      handleAnalyzeText(message.payload as AnalyzePayload)
        .then(sendResponse)
        .catch((err) => sendResponse({ error: String(err) }));
      return true; // Keep channel open for async response
    }
  });
});