docfast/src/services/keys.ts

import { randomBytes } from "crypto";
import logger from "./logger.js";
import pool from "./db.js";
import { queryWithRetry, connectWithRetry } from "./db.js";

export interface ApiKey {
  key: string;
  tier: "free" | "pro";
  email: string;
  createdAt: string;
  stripeCustomerId?: string;
}

// In-memory cache for fast lookups, synced with PostgreSQL
let keysCache: ApiKey[] = [];

export async function loadKeys(): Promise<void> {
  try {
    const result = await queryWithRetry(
      "SELECT key, tier, email, created_at, stripe_customer_id FROM api_keys"
    );
    keysCache = result.rows.map((r) => ({
      key: r.key,
      tier: r.tier as "free" | "pro",
      email: r.email,
      createdAt: r.created_at instanceof Date ? r.created_at.toISOString() : r.created_at,
      stripeCustomerId: r.stripe_customer_id || undefined,
    }));
  } catch (err) {
    logger.error({ err }, "Failed to load keys from PostgreSQL");
    keysCache = [];
  }

  // Also load seed keys from env
  const envKeys = process.env.API_KEYS?.split(",").map((k) => k.trim()).filter(Boolean) || [];
  for (const k of envKeys) {
    if (!keysCache.find((e) => e.key === k)) {
      const entry: ApiKey = { key: k, tier: "pro", email: "seed@docfast.dev", createdAt: new Date().toISOString() };
      keysCache.push(entry);
      // Upsert into DB
      await queryWithRetry(
        `INSERT INTO api_keys (key, tier, email, created_at) VALUES ($1, $2, $3, $4)
         ON CONFLICT (key) DO NOTHING`,
        [k, "pro", "seed@docfast.dev", new Date().toISOString()]
      ).catch(() => {});
    }
  }
}

export function isValidKey(key: string): boolean {
  return keysCache.some((k) => k.key === key);
}

export function getKeyInfo(key: string): ApiKey | undefined {
  return keysCache.find((k) => k.key === key);
}

export function isProKey(key: string): boolean {
  const info = getKeyInfo(key);
  return info?.tier === "pro";
}

function generateKey(prefix: string): string {
  return `${prefix}_${randomBytes(24).toString("hex")}`;
}

export async function createFreeKey(email?: string): Promise<ApiKey> {
  if (email) {
    const existing = keysCache.find((k) => k.email === email && k.tier === "free");
    if (existing) return existing;
  }

  const entry: ApiKey = {
    key: generateKey("df_free"),
    tier: "free",
    email: email || "",
    createdAt: new Date().toISOString(),
  };

  await queryWithRetry(
    "INSERT INTO api_keys (key, tier, email, created_at) VALUES ($1, $2, $3, $4)",
    [entry.key, entry.tier, entry.email, entry.createdAt]
  );
  keysCache.push(entry);
  return entry;
}

export async function createProKey(email: string, stripeCustomerId: string): Promise<ApiKey> {
  // Check in-memory cache first (fast path)
  const existing = keysCache.find((k) => k.stripeCustomerId === stripeCustomerId);
  if (existing) {
    existing.tier = "pro";
    await queryWithRetry("UPDATE api_keys SET tier = 'pro' WHERE key = $1", [existing.key]);
    return existing;
  }

  // UPSERT: handles duplicate webhooks across pods via DB unique index
  const newKey = generateKey("df_pro");
  const now = new Date().toISOString();

  const result = await queryWithRetry(
    `INSERT INTO api_keys (key, tier, email, created_at, stripe_customer_id)
     VALUES ($1, $2, $3, $4, $5)
     ON CONFLICT (stripe_customer_id) WHERE stripe_customer_id IS NOT NULL
     DO UPDATE SET tier = 'pro'
     RETURNING key, tier, email, created_at, stripe_customer_id`,
    [newKey, "pro", email, now, stripeCustomerId]
  );

  const row = result.rows[0];
  const entry: ApiKey = {
    key: row.key,
    tier: row.tier as "free" | "pro",
    email: row.email,
    createdAt: row.created_at instanceof Date ? row.created_at.toISOString() : row.created_at,
    stripeCustomerId: row.stripe_customer_id || undefined,
  };

  // Refresh in-memory cache
  const cacheIdx = keysCache.findIndex((k) => k.stripeCustomerId === stripeCustomerId);
  if (cacheIdx >= 0) {
    keysCache[cacheIdx] = entry;
  } else {
    keysCache.push(entry);
  }

  return entry;
}

export async function downgradeByCustomer(stripeCustomerId: string): Promise<boolean> {
  const entry = keysCache.find((k) => k.stripeCustomerId === stripeCustomerId);
  if (entry) {
    entry.tier = "free";
    await queryWithRetry("UPDATE api_keys SET tier = 'free' WHERE stripe_customer_id = $1", [stripeCustomerId]);
    return true;
  }

  // DB fallback: key may exist on another pod's cache or after a restart
  logger.info({ stripeCustomerId }, "downgradeByCustomer: cache miss, falling back to DB");
  const result = await queryWithRetry(
    "SELECT key, tier, email, created_at, stripe_customer_id FROM api_keys WHERE stripe_customer_id = $1 LIMIT 1",
    [stripeCustomerId]
  );
  if (result.rows.length === 0) {
    logger.warn({ stripeCustomerId }, "downgradeByCustomer: customer not found in cache or DB");
    return false;
  }

  const row = result.rows[0];
  await queryWithRetry("UPDATE api_keys SET tier = 'free' WHERE stripe_customer_id = $1", [stripeCustomerId]);

  // Add to local cache so subsequent lookups on this pod work
  const cached: ApiKey = {
    key: row.key,
    tier: "free",
    email: row.email,
    createdAt: row.created_at instanceof Date ? row.created_at.toISOString() : row.created_at,
    stripeCustomerId: row.stripe_customer_id || undefined,
  };
  keysCache.push(cached);

  logger.info({ stripeCustomerId, key: row.key }, "downgradeByCustomer: downgraded via DB fallback");
  return true;
}

export async function findKeyByCustomerId(stripeCustomerId: string): Promise<ApiKey | null> {
  // Check DB directly — survives pod restarts unlike in-memory cache
  const result = await queryWithRetry(
    "SELECT key, tier, email, created_at, stripe_customer_id FROM api_keys WHERE stripe_customer_id = $1 LIMIT 1",
    [stripeCustomerId]
  );
  if (result.rows.length === 0) return null;
  const r = result.rows[0];
  return {
    key: r.key,
    tier: r.tier as "free" | "pro",
    email: r.email,
    createdAt: r.created_at instanceof Date ? r.created_at.toISOString() : r.created_at,
    stripeCustomerId: r.stripe_customer_id || undefined,
  };
}

export function getAllKeys(): ApiKey[] {
  return [...keysCache];
}

export async function updateKeyEmail(apiKey: string, newEmail: string): Promise<boolean> {
  const entry = keysCache.find((k) => k.key === apiKey);
  if (entry) {
    entry.email = newEmail;
    await queryWithRetry("UPDATE api_keys SET email = $1 WHERE key = $2", [newEmail, apiKey]);
    return true;
  }

  // DB fallback: key may exist on another pod's cache or after a restart
  logger.info({ apiKey: apiKey.slice(0, 10) + "..." }, "updateKeyEmail: cache miss, falling back to DB");
  const result = await queryWithRetry(
    "SELECT key, tier, email, created_at, stripe_customer_id FROM api_keys WHERE key = $1 LIMIT 1",
    [apiKey]
  );
  if (result.rows.length === 0) {
    logger.warn({ apiKey: apiKey.slice(0, 10) + "..." }, "updateKeyEmail: key not found in cache or DB");
    return false;
  }

  const row = result.rows[0];
  await queryWithRetry("UPDATE api_keys SET email = $1 WHERE key = $2", [newEmail, apiKey]);

  // Hydrate local cache
  const cached: ApiKey = {
    key: row.key,
    tier: row.tier as "free" | "pro",
    email: newEmail,
    createdAt: row.created_at instanceof Date ? row.created_at.toISOString() : row.created_at,
    stripeCustomerId: row.stripe_customer_id || undefined,
  };
  keysCache.push(cached);

  logger.info({ apiKey: apiKey.slice(0, 10) + "..." }, "updateKeyEmail: updated via DB fallback");
  return true;
}

export async function updateEmailByCustomer(stripeCustomerId: string, newEmail: string): Promise<boolean> {
  const entry = keysCache.find(k => k.stripeCustomerId === stripeCustomerId);
  if (entry) {
    entry.email = newEmail;
    await queryWithRetry("UPDATE api_keys SET email = $1 WHERE stripe_customer_id = $2", [newEmail, stripeCustomerId]);
    return true;
  }

  // DB fallback: key may exist on another pod's cache or after a restart
  logger.info({ stripeCustomerId }, "updateEmailByCustomer: cache miss, falling back to DB");
  const result = await queryWithRetry(
    "SELECT key, tier, email, created_at, stripe_customer_id FROM api_keys WHERE stripe_customer_id = $1 LIMIT 1",
    [stripeCustomerId]
  );
  if (result.rows.length === 0) {
    logger.warn({ stripeCustomerId }, "updateEmailByCustomer: customer not found in cache or DB");
    return false;
  }

  const row = result.rows[0];
  await queryWithRetry("UPDATE api_keys SET email = $1 WHERE stripe_customer_id = $2", [newEmail, stripeCustomerId]);

  // Hydrate local cache
  const cached: ApiKey = {
    key: row.key,
    tier: row.tier as "free" | "pro",
    email: newEmail,
    createdAt: row.created_at instanceof Date ? row.created_at.toISOString() : row.created_at,
    stripeCustomerId: row.stripe_customer_id || undefined,
  };
  keysCache.push(cached);

  logger.info({ stripeCustomerId, key: row.key }, "updateEmailByCustomer: updated via DB fallback");
  return true;
}