server/webSearchUsageStore.js

import fs from 'fs/promises';
import path from 'path';
import { DATA_ROOT, isPostgresStorageMode } from './dataPaths.js';
import {
  decryptJsonPayload,
  encryptJsonPayload,
  makeLookupToken,
  pgQuery,
  withPgTransaction,
} from './postgres.js';

const STORE_FILE = path.join(DATA_ROOT, 'web-search-usage.json');
const APP_TIMEZONE = process.env.APP_TIMEZONE || 'America/New_York';

let state = { days: {} };
let loaded = false;
let saveChain = Promise.resolve();

function todayKey() {
  return new Intl.DateTimeFormat('en-CA', {
    timeZone: APP_TIMEZONE,
    year: 'numeric',
    month: '2-digit',
    day: '2-digit',
  }).format(new Date());
}

function usageLookup(key) {
  return makeLookupToken('web-search-usage', key);
}

function usageAad(key, day) {
  return `web-search-usage:${key}:${day}`;
}

function pruneDays() {
  const keepKey = todayKey();
  state.days = Object.fromEntries(
    Object.entries(state.days || {}).filter(([key]) => key === keepKey)
  );
}

async function ensureLoaded() {
  if (loaded || isPostgresStorageMode()) return;
  loaded = true;
  try {
    await fs.mkdir(DATA_ROOT, { recursive: true });
    const raw = await fs.readFile(STORE_FILE, 'utf8');
    const parsed = JSON.parse(raw);
    if (parsed && typeof parsed === 'object') state = parsed;
  } catch {}
  pruneDays();
}

function saveState() {
  saveChain = saveChain.then(async () => {
    pruneDays();
    await fs.mkdir(DATA_ROOT, { recursive: true });
    await fs.writeFile(STORE_FILE, JSON.stringify(state, null, 2), 'utf8');
  }).catch((err) => {
    console.error('Failed to persist web search usage:', err);
  });
  return saveChain;
}

function getCounterRecord(key) {
  const day = todayKey();
  if (!state.days[day]) state.days[day] = {};
  if (!state.days[day][key]) state.days[day][key] = 0;
  return {
    day,
    used: state.days[day][key],
    set(nextValue) {
      state.days[day][key] = nextValue;
    },
  };
}

async function pruneSqlUsage(day, client = null) {
  const runner = client || { query: pgQuery };
  await runner.query('DELETE FROM web_search_usage WHERE day_key <> $1', [day]);
}

async function getSqlUsage(key, limit = 15) {
  const day = todayKey();
  await pruneSqlUsage(day);
  const lookup = usageLookup(key);
  const { rows } = await pgQuery(
    'SELECT payload FROM web_search_usage WHERE key_lookup = $1 AND day_key = $2',
    [lookup, day]
  );
  const payload = rows[0]
    ? decryptJsonPayload(rows[0].payload, usageAad(key, day))
    : null;
  const used = Math.max(0, Number(payload?.used) || 0);
  return {
    limit,
    used,
    remaining: Math.max(0, limit - used),
    window: day,
    period: 'daily',
  };
}

async function consumeSqlUsage(key, limit = 15) {
  return withPgTransaction(async (client) => {
    const day = todayKey();
    await pruneSqlUsage(day, client);
    const lookup = usageLookup(key);
    const { rows } = await client.query(
      'SELECT payload FROM web_search_usage WHERE key_lookup = $1 AND day_key = $2 FOR UPDATE',
      [lookup, day]
    );
    const current = rows[0]
      ? decryptJsonPayload(rows[0].payload, usageAad(key, day))
      : { used: 0 };
    const used = Math.max(0, Number(current?.used) || 0);

    if (used >= limit) {
      return {
        allowed: false,
        limit,
        used,
        remaining: 0,
        window: day,
        period: 'daily',
      };
    }

    const next = { used: used + 1 };
    await client.query(
      `INSERT INTO web_search_usage (key_lookup, day_key, updated_at, payload)
       VALUES ($1, $2, $3, $4::jsonb)
       ON CONFLICT (key_lookup, day_key)
       DO UPDATE SET updated_at = EXCLUDED.updated_at, payload = EXCLUDED.payload`,
      [lookup, day, new Date().toISOString(), JSON.stringify(encryptJsonPayload(next, usageAad(key, day)))]
    );

    return {
      allowed: true,
      limit,
      used: next.used,
      remaining: Math.max(0, limit - next.used),
      window: day,
      period: 'daily',
    };
  });
}

export async function getWebSearchUsage(key, limit = 15) {
  if (isPostgresStorageMode()) return getSqlUsage(key, limit);

  await ensureLoaded();
  const record = getCounterRecord(key);
  return {
    limit,
    used: record.used,
    remaining: Math.max(0, limit - record.used),
    window: record.day,
    period: 'daily',
  };
}

export async function consumeWebSearchUsage(key, limit = 15) {
  if (isPostgresStorageMode()) return consumeSqlUsage(key, limit);

  await ensureLoaded();
  const record = getCounterRecord(key);
  if (record.used >= limit) {
    return {
      allowed: false,
      ...(await getWebSearchUsage(key, limit)),
    };
  }
  record.set(record.used + 1);
  await saveState();
  return {
    allowed: true,
    ...(await getWebSearchUsage(key, limit)),
  };
}