Epicure/apps/web/lib/ai/resolve-user-key.ts

import { db, userAiKeys, userModelPrefs, eq, and } from "@epicure/db";
import { decrypt } from "@/lib/encrypt";
import { getSiteSetting, type SiteSettingKey } from "@/lib/site-settings";
import type { AiConfig, AiProvider } from "./factory";

export type ModelUseCase = "text" | "vision" | "mealPlan";

const USE_CASE_SITE_DEFAULTS: Record<ModelUseCase, { provider: SiteSettingKey; model: SiteSettingKey }> = {
  text: { provider: "DEFAULT_TEXT_PROVIDER", model: "DEFAULT_TEXT_MODEL" },
  vision: { provider: "DEFAULT_VISION_PROVIDER", model: "DEFAULT_VISION_MODEL" },
  mealPlan: { provider: "DEFAULT_MEAL_PLAN_PROVIDER", model: "DEFAULT_MEAL_PLAN_MODEL" },
};

const PROVIDER_API_KEY_SETTING: Record<AiProvider, SiteSettingKey | null> = {
  openrouter: "OPENROUTER_API_KEY",
  openai: "OPENAI_API_KEY",
  anthropic: "ANTHROPIC_API_KEY",
  ollama: null,
};

/**
 * Thrown when a user's saved BYOK key exists but fails to decrypt (e.g. the
 * encryption secret rotated, or the stored ciphertext is corrupt). Callers
 * must surface this to the user instead of silently falling back to the
 * platform key — otherwise the user believes their own key/billing is being
 * used when it is not.
 */
export class ByokDecryptError extends Error {
  constructor(public readonly provider: string) {
    super(`Failed to decrypt saved API key for provider "${provider}". Please re-enter it in Settings.`);
    this.name = "ByokDecryptError";
  }
}

export async function withUserKey(userId: string, config: AiConfig): Promise<AiConfig> {
  const provider = config.provider;
  if (!provider) return config;

  const row = await db.query.userAiKeys.findFirst({
    where: and(eq(userAiKeys.userId, userId), eq(userAiKeys.provider, provider)),
  });

  if (!row) return config;

  try {
    const apiKey = decrypt(row.encryptedKey);
    return { ...config, apiKey };
  } catch {
    throw new ByokDecryptError(provider);
  }
}

export async function getModelConfigForUseCase(userId: string, useCase: ModelUseCase): Promise<AiConfig> {
  const prefs = await db.query.userModelPrefs.findFirst({
    where: eq(userModelPrefs.userId, userId),
  });

  const providerField = `${useCase}Provider` as keyof typeof prefs;
  const modelField = `${useCase}Model` as keyof typeof prefs;
  const provider = prefs?.[providerField] as AiProvider | null | undefined;
  const model = prefs?.[modelField] as string | null | undefined;

  if (provider) {
    const config: AiConfig = { provider, model: model ?? undefined };
    return withUserKey(userId, config);
  }

  return getDefaultProviderWithKey(userId, useCase);
}

export async function getDefaultProviderWithKey(userId: string, useCase?: ModelUseCase): Promise<AiConfig> {
  const keys = await db.query.userAiKeys.findMany({
    where: eq(userAiKeys.userId, userId),
    columns: { provider: true, encryptedKey: true },
  });

  const order: AiProvider[] = ["openrouter", "openai", "anthropic", "ollama"];

  // 1. User's own BYOK key takes highest priority — it's their own credentials/billing.
  for (const p of order) {
    const row = keys.find((k) => k.provider === p);
    if (row) {
      try {
        return { provider: p, apiKey: decrypt(row.encryptedKey) };
      } catch {
        throw new ByokDecryptError(p);
      }
    }
  }

  // 2. Admin-configured default provider/model for this use case, if set.
  if (useCase) {
    const setting = USE_CASE_SITE_DEFAULTS[useCase];
    const defaultProvider = (await getSiteSetting(setting.provider)) as AiProvider | null;
    if (defaultProvider) {
      const defaultModel = await getSiteSetting(setting.model);
      const apiKeySetting = PROVIDER_API_KEY_SETTING[defaultProvider];
      const apiKey = apiKeySetting ? await getSiteSetting(apiKeySetting) : null;
      return { provider: defaultProvider, model: defaultModel ?? undefined, apiKey: apiKey ?? undefined };
    }
  }

  // 3. Admin site-setting API keys, first provider found (no explicit default set)
  for (const p of order) {
    const apiKeySetting = PROVIDER_API_KEY_SETTING[p];
    if (!apiKeySetting) continue;
    const apiKey = await getSiteSetting(apiKeySetting);
    if (apiKey) return { provider: p, apiKey };
  }

  // 4. Let factory use raw env vars
  return {};
}