lila/data-pipeline/staging.ts

import { randomUUID } from "node:crypto";
import { readFileSync } from "node:fs";
import Database from "better-sqlite3";
import type { SupportedLanguageCode, SupportedPos } from "@lila/shared";
import type { GeminiWordEntry } from "./validate.js";

const STAGING_TABLES = ["words", "senses", "translations"] as const;

export const openStaging = (
  dbPath: string,
  schemaPath: string,
): Database.Database => {
  const db = new Database(dbPath);
  db.pragma("foreign_keys = ON");

  const existing = db
    .prepare<[], { name: string }>(
      `SELECT name FROM sqlite_master WHERE type = 'table' AND name IN ('words', 'senses', 'translations')`,
    )
    .all()
    .map((row) => row.name);

  if (existing.length === 0) {
    db.exec(readFileSync(schemaPath, "utf-8"));
  } else if (existing.length !== STAGING_TABLES.length) {
    db.close();
    throw new Error(
      `staging database at ${dbPath} has a partial schema (found: ${existing.join(", ")}) — fix or delete it`,
    );
  }
  return db;
};

export const getStagedHeadwords = (
  db: Database.Database,
  language: SupportedLanguageCode,
  pos: SupportedPos,
): Set<string> => {
  const rows = db
    .prepare<
      [string, string],
      { headword: string }
    >(`SELECT headword FROM words WHERE language_code = ? AND pos = ?`)
    .all(language, pos);
  return new Set(rows.map((row) => row.headword));
};

/**
 * Writes one validated entry (word + senses + translations) atomically.
 * Returns "already-staged" without writing if the word exists — a partial
 * word can never be left behind, so no NOT NULL relaxation is needed.
 */
export const stageEntry = (
  db: Database.Database,
  entry: GeminiWordEntry,
): "staged" | "already-staged" => {
  const insert = db.transaction((): "staged" | "already-staged" => {
    const existing = db
      .prepare<
        [string, string, string],
        { id: string }
      >(`SELECT id FROM words WHERE headword = ? AND language_code = ? AND pos = ?`)
      .get(entry.headword, entry.language, entry.pos);
    if (existing !== undefined) {
      return "already-staged";
    }

    const wordId = randomUUID();
    db.prepare(
      `INSERT INTO words (id, headword, language_code, pos) VALUES (?, ?, ?, ?)`,
    ).run(wordId, entry.headword, entry.language, entry.pos);

    const insertSense = db.prepare(
      `INSERT INTO senses (id, word_id, sense_index, difficulty, definitions, examples)
       VALUES (?, ?, ?, ?, ?, ?)`,
    );
    const insertTranslation = db.prepare(
      `INSERT INTO translations (id, sense_id, target_language_code, translation, gender, difficulty)
       VALUES (?, ?, ?, ?, ?, ?)`,
    );

    for (const sense of entry.senses) {
      const senseId = randomUUID();
      insertSense.run(
        senseId,
        wordId,
        sense.sense_index,
        sense.difficulty,
        JSON.stringify(sense.definitions),
        JSON.stringify(sense.examples),
      );
      for (const translation of sense.translations) {
        insertTranslation.run(
          randomUUID(),
          senseId,
          translation.target_language,
          translation.word,
          translation.gender,
          translation.difficulty,
        );
      }
    }
    return "staged";
  });
  return insert();
};

export type StagingCounts = {
  words: number;
  senses: number;
  translations: number;
};

export const countStagedRows = (db: Database.Database): StagingCounts => {
  const count = (table: (typeof STAGING_TABLES)[number]): number =>
    db.prepare<[], { n: number }>(`SELECT COUNT(*) AS n FROM ${table}`).get()
      ?.n ?? 0;
  return {
    words: count("words"),
    senses: count("senses"),
    translations: count("translations"),
  };
};