LexAI/src/lib/providers.ts

// Provider adapter table + the single LLM call path. Replaces the eight
// hand-rolled per-provider functions (callX / callXWithPrompt families) that
// previously lived in the background worker. Behavior-preserving: URLs,
// headers, request bodies, default models, and error strings match the old
// implementations — pinned by tests/unit/providers.test.ts.

import type { LexAIConfig, LexAIResponse, PromptParams } from './types';

export async function fetchWithTimeout(
  url: string,
  options: RequestInit,
  timeoutMs = 30000,
): Promise<Response> {
  const controller = new AbortController();
  const id = setTimeout(() => controller.abort(), timeoutMs);
  try {
    return await fetch(url, { ...options, signal: controller.signal });
  } finally {
    clearTimeout(id);
  }
}

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

// ─── Prompt Builder pattern library (docs/prompting_style.md) ────────────────
// Composed in a fixed order — ROLE+TASK, PATTERN, MODIFIERS, INVARIANTS last
// (invariants last so they survive context decay, guide §1) — never the whole
// table at once: only the selected pattern's instruction+skeleton+guard.

const PROMPT_BASE_TASK =
  'You are an expert prompt engineer. Transform the provided text into one well-crafted prompt that will get the best possible result from an AI model. ' +
  'First infer what the user is trying to achieve — the text may be a rough idea, a question, or a description of the output they want — then compose exactly one engineered prompt for that goal.';

// Auto routing rubric — lives in the base so 'auto' (or an absent/unknown
// pattern) adds nothing beyond it. Mirrors the guide's summary matrix.
const PROMPT_AUTO_RUBRIC =
  ' Choose the structural pattern the goal actually needs: formatting or extraction favors few-shot examples or an output contract; ' +
  'logic or multi-constraint problems favor chain-of-thought; open-ended design favors plan-and-solve or tree-of-thoughts; ' +
  'tasks needing live data or tools favor ReAct; multi-step builds or migrations favor plan-execute-verify; ' +
  'tasks that must beat a quality bar favor a gauntlet builder-judge loop; otherwise use a plain zero-shot instruction. ' +
  'Prefer the cheapest pattern that meets the goal — never add reasoning scaffolding to a simple task.';

// One instruction+skeleton+failure-mode-guard per PROMPT_PATTERNS id
// (excluding 'auto', which uses the rubric above instead).
const PROMPT_PATTERN_INSTRUCTIONS: Record<string, string> = {
  'zero-shot':
    ' Compose the engineered prompt as direct, imperative instructions — one clear task per sentence, ' +
    'with no scaffolding beyond what the task actually needs.',
  role:
    ' Compose the engineered prompt as a system-role block with ROLE, DOMAIN, INVARIANT, and OUTPUT_FORMAT lines that assign the model ' +
    'its expertise and constraints; restate the single hardest constraint again as the very last line, to guard against it being forgotten ' +
    'in a long conversation (context decay).',
  'few-shot':
    ' Compose the engineered prompt using explicit <example><input>…</input><output>…</output></example> delimiters, with 1-2 examples ' +
    'that are structurally diverse from each other (not near-duplicates), to guard against the model overfitting to the last example\'s ' +
    'exact values (recency/label bias), then a trailing open <target><input>…</input><output> for the real input.',
  structured:
    ' Compose the engineered prompt as labeled sections in this order: Context, Task, Constraints, Output format — each a short heading ' +
    'followed by its content.',
  contract:
    ' Compose the engineered prompt around an exact output schema: name every field, its type, and whether it is required, and include a ' +
    'rule instructing the model to reject or omit anything outside that schema.',
  cot:
    ' Compose the engineered prompt with two explicit phases labeled PHASE 1: REASONING and PHASE 2: OUTPUT — the model works through its ' +
    'reasoning in phase 1, then gives a final answer in phase 2 that stands alone without needing the reasoning to make sense.',
  'plan-solve':
    ' Compose the engineered prompt so the model must first produce a numbered plan of the steps needed, then execute that plan in order, ' +
    'referencing each step as it completes it.',
  tot:
    ' Compose the engineered prompt instructing the model to generate several candidate approaches, score each 0-1 against stated criteria, ' +
    'prune the weak ones, expand on the best, and report the winning approach and why it was chosen.',
  react:
    ' Compose the engineered prompt for a tool-capable agent, not a plain chat model: declare the available tools, require a strict ' +
    'Thought: / Action: / Observation: loop, and terminate with a line starting Final Answer:. Include a hard step budget (e.g. max 10 steps) ' +
    'and a rule that repeating the same action signature twice must break the loop, to guard against infinite loops.',
  pev:
    ' Compose the engineered prompt for a tool-capable agent, not a plain chat model: require it to first generate a task DAG of sub-tasks, ' +
    'run a verification assertion after each node, re-plan the remaining DAG on any failed assertion, and state an explicit stop condition ' +
    'for when the task is complete.',
  gauntlet:
    ' Compose the engineered prompt for a tool-capable agent, not a plain chat model, running a builder-judge loop: the builder produces an ' +
    'artifact, a judge instantiated in a fresh context compares it against a NAMED reference standard, and returns exactly ' +
    'STATUS: [PASS|FAIL] | FEEDBACK: <one directive>; the loop stops on PASS or after a stated maximum number of rounds.',
};

function promptPatternSection(pattern?: string): string {
  if (pattern && pattern !== 'auto' && PROMPT_PATTERN_INSTRUCTIONS[pattern]) {
    return PROMPT_PATTERN_INSTRUCTIONS[pattern];
  }
  return PROMPT_AUTO_RUBRIC;
}

// Persona + output-format modifiers (Prompt Builder parameters). 'Auto'/absent
// adds nothing for either — the base prompt already tells the engineer to
// decide these from the input.
function promptParamModifiers(params?: PromptParams): string {
  if (!params) return '';
  const parts: string[] = [];

  if (params.persona === 'None') {
    parts.push(' Do not assign a persona or role in the engineered prompt.');
  } else if (params.persona && params.persona !== 'Auto') {
    parts.push(` The engineered prompt must assign the model the persona of ${params.persona.trim()}, including the skills and expertise that persona implies.`);
  }

  if (params.format && params.format !== 'Auto') {
    parts.push(` The engineered prompt must require the final output as ${params.format.toLowerCase()}.`);
  }

  return parts.join('');
}

// Hard invariants — always last, so they survive context decay on long
// compositions (guide §1: restate/keep critical rules at the end).
const PROMPT_INVARIANTS =
  ' Return ONLY the engineered prompt, ready to paste into an AI chat — no explanations, no surrounding quotes, no preamble. ' +
  'Keep it self-contained, and use the cheapest structure that meets the goal. ' +
  'If the input is missing information the prompt needs, mark it as a [BRACKETED] placeholder rather than inventing facts.';

export function getSystemPrompt(action: string, style?: string, promptParams?: PromptParams): string {
  // Normalize 'fix' (used by context menu and popup) to 'grammar'
  const normalizedAction = action === 'fix' ? 'grammar' : action;
  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.',
    explain:
      'You are a helpful teacher. Explain the following text in simple, easy-to-understand language. ' +
      'Break down complex terms, jargon, or concepts so anyone can understand. ' +
      'Be concise but clear. Return only the explanation, no extra commentary.',
    prompt: PROMPT_BASE_TASK,
  };
  const base = prompts[normalizedAction] ?? prompts.grammar;

  if (normalizedAction !== 'prompt') {
    // For every other action, style describes the wording of OUR response.
    const styleModifier = style && style !== 'Default' ? ` Write in a ${style.toLowerCase()} style.` : '';
    return base + styleModifier;
  }

  // 'prompt' action: ROLE+TASK, then PATTERN, then MODIFIERS (persona,
  // format, response style — in that order), then hard INVARIANTS last.
  const patternSection = promptPatternSection(promptParams?.pattern);
  const modifiers = promptParamModifiers(promptParams);
  // The style here describes the output the ENGINEERED prompt should ask
  // for — not the wording of this system prompt.
  const styleModifier = style && style !== 'Default'
    ? ` The engineered prompt should instruct the model to respond in a ${style.toLowerCase()} style.`
    : '';
  return base + patternSection + modifiers + styleModifier + PROMPT_INVARIANTS;
}

// ─── Adapter table ────────────────────────────────────────────────────────────

export interface ProviderSpec {
  label: string; // human name used in error messages ('OpenAI error: …')
  chatUrl: string;
  modelsUrl: string;
  defaultModel: string;
  headers: (apiKey: string) => Record<string, string>;
  body: (model: string, systemPrompt: string, text: string, maxTokens: number) => Record<string, unknown>;
  extract: (data: any) => string | undefined;
}

function bearerHeaders(apiKey: string): Record<string, string> {
  return { 'Content-Type': 'application/json', Authorization: `Bearer ${apiKey}` };
}

// OpenAI-compatible chat body (OpenAI, Groq, OpenRouter). OpenRouter's old
// implementation sent no temperature — preserve that.
function openAiStyleBody(temperature?: number) {
  return (model: string, systemPrompt: string, text: string, maxTokens: number) => ({
    model,
    messages: [
      { role: 'system', content: systemPrompt },
      { role: 'user', content: text },
    ],
    max_tokens: maxTokens,
    ...(temperature !== undefined ? { temperature } : {}),
  });
}

const extractOpenAiStyle = (data: any): string | undefined => data?.choices?.[0]?.message?.content;

// OpenAI proper rejects the legacy `max_tokens` on newer models ("Use
// 'max_completion_tokens' instead") and reasoning models (o-series, gpt-5
// family) additionally reject any temperature other than the default.
// `max_completion_tokens` is accepted by all current OpenAI chat models, so
// send it unconditionally; temperature stays only for non-reasoning models.
// Groq/OpenRouter still expect `max_tokens` — don't apply this to them.
const OPENAI_REASONING_MODEL_RE = /^(o\d|gpt-5)/;

function openAiBody(model: string, systemPrompt: string, text: string, maxTokens: number) {
  return {
    model,
    messages: [
      { role: 'system', content: systemPrompt },
      { role: 'user', content: text },
    ],
    max_completion_tokens: maxTokens,
    ...(OPENAI_REASONING_MODEL_RE.test(model) ? {} : { temperature: 0.7 }),
  };
}

export const PROVIDER_SPECS: Record<string, ProviderSpec> = {
  openai: {
    label: 'OpenAI',
    chatUrl: 'https://api.openai.com/v1/chat/completions',
    modelsUrl: 'https://api.openai.com/v1/models',
    defaultModel: 'gpt-4o-mini',
    headers: bearerHeaders,
    body: openAiBody,
    extract: extractOpenAiStyle,
  },
  anthropic: {
    label: 'Anthropic',
    chatUrl: 'https://api.anthropic.com/v1/messages',
    modelsUrl: 'https://api.anthropic.com/v1/models',
    defaultModel: 'claude-3-5-haiku-20241022',
    headers: (apiKey) => ({
      'Content-Type': 'application/json',
      'x-api-key': apiKey,
      'anthropic-version': '2023-06-01',
      // Anthropic rejects browser-origin requests unless this opt-in is sent.
      // The service worker counts as a browser origin, so it's required here too.
      'anthropic-dangerous-direct-browser-access': 'true',
    }),
    body: (model, systemPrompt, text, maxTokens) => ({
      model,
      max_tokens: maxTokens,
      system: systemPrompt,
      messages: [{ role: 'user', content: text }],
    }),
    extract: (data) => data?.content?.[0]?.text,
  },
  groq: {
    label: 'Groq',
    chatUrl: 'https://api.groq.com/openai/v1/chat/completions',
    modelsUrl: 'https://api.groq.com/openai/v1/models',
    defaultModel: 'llama-3.3-70b-versatile',
    headers: bearerHeaders,
    body: openAiStyleBody(0.7),
    extract: extractOpenAiStyle,
  },
  openrouter: {
    label: 'OpenRouter',
    chatUrl: 'https://openrouter.ai/api/v1/chat/completions',
    modelsUrl: 'https://openrouter.ai/api/v1/models',
    defaultModel: 'openai/gpt-4o-mini',
    headers: (apiKey) => ({
      ...bearerHeaders(apiKey),
      'HTTP-Referer': 'https://lexai.dev',
      'X-Title': 'LexAI',
    }),
    body: openAiStyleBody(undefined),
    extract: extractOpenAiStyle,
  },
};

// Display name for a provider id ('groq' → 'Groq'), for user-facing messages.
export function providerLabel(provider: string): string {
  return PROVIDER_SPECS[provider]?.label ?? provider;
}

// A rejected key is the one failure users can actually fix, and the provider's
// own wording ("Invalid API Key") doesn't say where to fix it.
const KEY_HINT = ' — open LexAI Settings and re-enter your API key for this provider.';
const isAuthStatus = (status: number) => status === 401 || status === 403;

// ─── Chat call ────────────────────────────────────────────────────────────────

// Scale the output budget with the input instead of the old hard-coded 1024
// (which truncated "Expand" on long selections). chars ≈ tokens × 4, so this
// allows roughly 4× the input length in output, clamped to a sane range.
export function defaultMaxTokens(text: string): number {
  return Math.max(1024, Math.min(8192, Math.ceil(text.length)));
}

export interface CallOptions {
  maxTokens?: number;
}

export async function callProvider(
  config: LexAIConfig,
  text: string,
  systemPrompt: string,
  opts?: CallOptions,
): Promise<LexAIResponse> {
  const provider = config.provider || 'openai';
  const spec = PROVIDER_SPECS[provider];
  if (!spec) {
    return { error: `Unknown provider: "${provider}". Please check LexAI settings.` };
  }

  const model = config.model || spec.defaultModel;
  const maxTokens = opts?.maxTokens ?? defaultMaxTokens(text);

  let res: Response;
  try {
    res = await fetchWithTimeout(spec.chatUrl, {
      method: 'POST',
      headers: spec.headers(config.apiKey ?? ''),
      body: JSON.stringify(spec.body(model, systemPrompt, text, maxTokens)),
    });
  } catch (err) {
    return { error: `Network error reaching ${spec.label}: ${String(err)}` };
  }

  // Guard the parse — gateways/proxies return HTML error pages (502 etc.)
  // which would otherwise surface as a raw SyntaxError to the user.
  const data = await res.json().catch(() => null);
  if (!res.ok) {
    const msg = data?.error?.message ?? `HTTP ${res.status}`;
    return { error: `${spec.label} error: ${msg}${isAuthStatus(res.status) ? KEY_HINT : ''}` };
  }

  const result = spec.extract(data);
  if (!result) return { error: `${spec.label} returned an empty response.` };
  return { result: result.trim() };
}

// ─── Live model listing ───────────────────────────────────────────────────────

// Drop non-chat models (embeddings, audio, image, etc.) so the picker stays useful.
const NON_CHAT_MODEL_RE = /embedding|whisper|tts|dall-e|audio|realtime|moderation|image|guard|transcribe|speech|rerank/i;

export async function listModels(
  provider: string,
  apiKey?: string,
): Promise<{ models?: string[]; error?: string; keyRejected?: boolean }> {
  const spec = PROVIDER_SPECS[provider];
  if (!spec) return { error: `Unknown provider: "${provider}". Please check LexAI settings.` };

  const headers: Record<string, string> = { 'Content-Type': 'application/json' };
  if (provider === 'anthropic') {
    if (!apiKey) return { error: 'Anthropic requires an API key to list models.' };
    headers['x-api-key'] = apiKey;
    headers['anthropic-version'] = '2023-06-01';
    // Allow the extension origin to call Anthropic directly (avoids a CORS 403).
    headers['anthropic-dangerous-direct-browser-access'] = 'true';
  } else if (apiKey) {
    // OpenRouter's list is public, so the key is optional there; OpenAI/Groq require it.
    headers['Authorization'] = `Bearer ${apiKey}`;
  }

  let res: Response;
  try {
    res = await fetchWithTimeout(spec.modelsUrl, { method: 'GET', headers }, 15000);
  } catch (err) {
    return { error: `Network error reaching ${spec.label}: ${String(err)}` };
  }

  const data = await res.json().catch(() => null);
  if (!res.ok) {
    const msg = data?.error?.message ?? data?.error ?? `HTTP ${res.status}`;
    // Flag rejected keys so Options can tell the user to re-enter one instead
    // of leaving the 🔒 badge implying the stored key is usable.
    return { error: `${spec.label} error: ${msg}`, ...(isAuthStatus(res.status) ? { keyRejected: true } : {}) };
  }

  const raw = Array.isArray(data?.data) ? data.data : [];
  const ids = raw
    .map((m: any) => (typeof m === 'string' ? m : m?.id))
    .filter((id: unknown): id is string => typeof id === 'string' && id.length > 0)
    .filter((id: string) => !NON_CHAT_MODEL_RE.test(id))
    .sort((a: string, b: string) => a.localeCompare(b));

  if (ids.length === 0) return { error: `No models returned by ${spec.label}.` };
  return { models: ids };
}