background.js

const OPENAI_URL = 'https://api.openai.com/v1/chat/completions';
const ANTHROPIC_URL = 'https://api.anthropic.com/v1/messages';
const OPENROUTER_URL = 'https://openrouter.ai/api/v1/chat/completions';

// Approximate pricing per 1M tokens { in, out } in USD
const MODEL_PRICING = {
  'gpt-4o':                    { in: 2.50,  out: 10.00 },
  'gpt-4o-mini':               { in: 0.15,  out: 0.60  },
  'gpt-4-turbo':               { in: 10.00, out: 30.00 },
  'gpt-4':                     { in: 30.00, out: 60.00 },
  'gpt-3.5-turbo':             { in: 0.50,  out: 1.50  },
  'o1':                        { in: 15.00, out: 60.00 },
  'o1-mini':                   { in: 3.00,  out: 12.00 },
  'o3-mini':                   { in: 1.10,  out: 4.40  },
  'claude-opus-4':             { in: 15.00, out: 75.00 },
  'claude-sonnet-4':           { in: 3.00,  out: 15.00 },
  'claude-3-5-sonnet':         { in: 3.00,  out: 15.00 },
  'claude-3-5-haiku':          { in: 0.80,  out: 4.00  },
  'claude-3-haiku':            { in: 0.25,  out: 1.25  },
  'claude-3-opus':             { in: 15.00, out: 75.00 },
  'claude-3-sonnet':           { in: 3.00,  out: 15.00 },
  'gemini-2.0-flash':          { in: 0.10,  out: 0.40  },
  'gemini-1.5-flash':          { in: 0.075, out: 0.30  },
  'gemini-1.5-pro':            { in: 1.25,  out: 5.00  },
  'llama-3.3-70b':             { in: 0.23,  out: 0.40  },
  'llama-3.1-8b':              { in: 0.06,  out: 0.06  },
  'deepseek-chat':             { in: 0.14,  out: 0.28  },
  'deepseek-r1':               { in: 0.55,  out: 2.19  },
  'deepseek-v4-flash':         { in: 0.14,  out: 0.28  },
  'mistral-7b':                { in: 0.07,  out: 0.07  },
  'mixtral-8x7b':              { in: 0.24,  out: 0.24  },
  'qwen-2.5-72b':              { in: 0.35,  out: 0.40  },
};

function estimateCost(model, inputTokens, outputTokens) {
  let p = MODEL_PRICING[model];
  if (!p) {
    // Try stripping OpenRouter prefix (e.g. "deepseek/deepseek-chat" → "deepseek-chat")
    const short = model.includes('/') ? model.split('/').pop() : model;
    p = MODEL_PRICING[short];
    if (!p) {
      // Prefix match — handles dated variants like 'gpt-4o-2024-11-20'
      for (const key of Object.keys(MODEL_PRICING)) {
        if (short.startsWith(key) || model.startsWith(key)) { p = MODEL_PRICING[key]; break; }
      }
    }
  }
  if (!p) return null; // unknown model — don't guess
  return (inputTokens * p.in + outputTokens * p.out) / 1_000_000;
}

async function recordUsage(provider, model, inputTokens, outputTokens) {
  if (!inputTokens && !outputTokens) return;
  const cost = estimateCost(model, inputTokens, outputTokens);
  return new Promise(resolve => {
    chrome.storage.local.get(['ww_usage'], data => {
      const prev = data.ww_usage || { requests: 0, input_tokens: 0, output_tokens: 0, est_cost_usd: 0, cost_known: false, since: Date.now() };
      const updated = {
        requests:      (prev.requests      || 0) + 1,
        input_tokens:  (prev.input_tokens  || 0) + inputTokens,
        output_tokens: (prev.output_tokens || 0) + outputTokens,
        est_cost_usd:  (prev.est_cost_usd  || 0) + (cost ?? 0),
        cost_known:    (prev.cost_known || false) || (cost !== null),
        since:         prev.since || Date.now(),
      };
      chrome.storage.local.set({ ww_usage: updated }, resolve);
    });
  });
}

chrome.runtime.onConnect.addListener((port) => {
  if (port.name !== 'ai-request') return;
  port.onMessage.addListener(async (message) => {
    try {
      const result = await handleAIRequest(message.payload);
      try { port.postMessage({ type: 'AI_RESULT', result }); } catch (_) {}
    } catch (err) {
      try { port.postMessage({ type: 'AI_ERROR', error: err.message || 'An unknown error occurred.' }); } catch (_) {}
    }
  });
});

async function handleAIRequest({ provider, apiKey, model, systemPrompt, userText }) {
  switch (provider) {
    case 'openai':     return callOpenAI(apiKey, model, systemPrompt, userText);
    case 'anthropic':  return callAnthropic(apiKey, model, systemPrompt, userText);
    case 'openrouter': return callOpenRouter(apiKey, model, systemPrompt, userText);
    default: throw new Error(`Unknown provider: ${provider}`);
  }
}

async function callOpenAI(apiKey, model, systemPrompt, userText) {
  const res = await fetch(OPENAI_URL, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json', 'Authorization': `Bearer ${apiKey}` },
    body: JSON.stringify({
      model,
      messages: [{ role: 'system', content: systemPrompt }, { role: 'user', content: userText }],
      temperature: 0.3,
      max_tokens: 2048,
    }),
  });
  const data = await res.json();
  if (!res.ok) throw new Error(data.error?.message || `OpenAI error ${res.status}`);
  const u = data.usage || {};
  await recordUsage('openai', model, u.prompt_tokens || 0, u.completion_tokens || 0);
  return { result: data.choices[0].message.content.trim(), usage: { input: u.prompt_tokens, output: u.completion_tokens } };
}

async function callAnthropic(apiKey, model, systemPrompt, userText) {
  const res = await fetch(ANTHROPIC_URL, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'x-api-key': apiKey,
      'anthropic-version': '2023-06-01',
    },
    body: JSON.stringify({
      model,
      max_tokens: 2048,
      system: systemPrompt,
      messages: [{ role: 'user', content: userText }],
    }),
  });
  const data = await res.json();
  if (!res.ok) throw new Error(data.error?.message || `Anthropic error ${res.status}`);
  const u = data.usage || {};
  await recordUsage('anthropic', model, u.input_tokens || 0, u.output_tokens || 0);
  return { result: data.content[0].text.trim(), usage: { input: u.input_tokens, output: u.output_tokens } };
}

async function callOpenRouter(apiKey, model, systemPrompt, userText) {
  const res = await fetch(OPENROUTER_URL, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Authorization': `Bearer ${apiKey}`,
      'HTTP-Referer': 'https://github.com/writewise-extension',
      'X-Title': 'WriteWise',
    },
    body: JSON.stringify({
      model,
      messages: [{ role: 'system', content: systemPrompt }, { role: 'user', content: userText }],
      temperature: 0.3,
    }),
  });
  const data = await res.json();
  if (!res.ok) throw new Error(data.error?.message || `OpenRouter error ${res.status}`);
  const u = data.usage || {};
  await recordUsage('openrouter', model, u.prompt_tokens || 0, u.completion_tokens || 0);
  return { result: data.choices[0].message.content.trim(), usage: { input: u.prompt_tokens, output: u.completion_tokens } };
}