import { PG_INSUFFICIENT_PRIVILEGE, PG_CONFIGURATION_LIMIT_EXCEEDED } from "@drdgvhbh/postgres-error-codes";
import { platform } from "node:os";
import { Writable } from "node:stream";
import { pipeline } from "node:stream/promises";
import postgres from "postgres";
import { must } from "../../../../../shared/src/must.js";
import { equals } from "../../../../../shared/src/set-utils.js";
import { createLiteTableStatement, createLiteIndexStatement } from "../../../db/create.js";
import { READONLY } from "../../../db/mode-enum.js";
import { TsvParser } from "../../../db/pg-copy.js";
import { mapPostgresToLite, mapPostgresToLiteIndex } from "../../../db/pg-to-lite.js";
import { getTypeParsers } from "../../../db/pg-type-parser.js";
import { runTx } from "../../../db/run-transaction.js";
import { importSnapshot, TransactionPool } from "../../../db/transaction-pool.js";
import { liteValue, JSON_STRINGIFIED } from "../../../types/lite.js";
import { liteTableName } from "../../../types/names.js";
import { pgClient } from "../../../types/pg.js";
import { CpuProfiler } from "../../../types/profiler.js";
import { ALLOWED_APP_ID_CHARACTERS } from "../../../types/shards.js";
import { id } from "../../../types/sql.js";
import { ReplicationStatusPublisher } from "../../replicator/replication-status.js";
import { ColumnMetadataStore } from "../../replicator/schema/column-metadata.js";
import { initReplicationState } from "../../replicator/schema/replication-state.js";
import { toStateVersionString } from "./lsn.js";
import { ensureShardSchema } from "./schema/init.js";
import { getPublicationInfo } from "./schema/published.js";
import { newReplicationSlot, replicationSlotExpression, validatePublications, addReplica, getInternalShardConfig, dropShard } from "./schema/shard.js";
async function initialSync(lc, shard, tx, upstreamURI, syncOptions, context) {
  if (!ALLOWED_APP_ID_CHARACTERS.test(shard.appID)) {
    throw new Error(
      "The App ID may only consist of lower-case letters, numbers, and the underscore character"
    );
  }
  const { tableCopyWorkers, profileCopy } = syncOptions;
  const copyProfiler = profileCopy ? await CpuProfiler.connect() : null;
  const sql = pgClient(lc, upstreamURI);
  const replicationSession = pgClient(lc, upstreamURI, {
    ["fetch_types"]: false,
    // Necessary for the streaming protocol
    connection: { replication: "database" }
    // https://www.postgresql.org/docs/current/protocol-replication.html
  });
  const slotName = newReplicationSlot(shard);
  const statusPublisher = new ReplicationStatusPublisher(tx).publish(
    lc,
    "Initializing"
  );
  try {
    await checkUpstreamConfig(sql);
    const { publications } = await ensurePublishedTables(lc, sql, shard);
    lc.info?.(`Upstream is setup with publications [${publications}]`);
    const { database, host } = sql.options;
    lc.info?.(`opening replication session to ${database}@${host}`);
    let slot;
    for (let first = true; ; first = false) {
      try {
        slot = await createReplicationSlot(lc, replicationSession, slotName);
        break;
      } catch (e) {
        if (first && e instanceof postgres.PostgresError) {
          if (e.code === PG_INSUFFICIENT_PRIVILEGE) {
            await sql`ALTER ROLE current_user WITH REPLICATION`;
            lc.info?.(`Added the REPLICATION role to database user`);
            continue;
          }
          if (e.code === PG_CONFIGURATION_LIMIT_EXCEEDED) {
            const slotExpression = replicationSlotExpression(shard);
            const dropped = await sql`
              SELECT slot_name as slot, pg_drop_replication_slot(slot_name) 
                FROM pg_replication_slots
                WHERE slot_name LIKE ${slotExpression} AND NOT active`;
            if (dropped.length) {
              lc.warn?.(
                `Dropped inactive replication slots: ${dropped.map(({ slot: slot2 }) => slot2)}`,
                e
              );
              continue;
            }
            lc.error?.(`Unable to drop replication slots`, e);
          }
        }
        throw e;
      }
    }
    const { snapshot_name: snapshot, consistent_point: lsn } = slot;
    const initialVersion = toStateVersionString(lsn);
    initReplicationState(tx, publications, initialVersion, context);
    const start = performance.now();
    const published = await runTx(
      sql,
      async (tx2) => {
        await tx2.unsafe(
          /* sql*/
          `SET TRANSACTION SNAPSHOT '${snapshot}'`
        );
        return getPublicationInfo(tx2, publications);
      },
      { mode: READONLY }
    );
    validatePublications(lc, published);
    const { tables, indexes } = published;
    const numTables = tables.length;
    if (platform() === "win32" && tableCopyWorkers < numTables) {
      lc.warn?.(
        `Increasing the number of copy workers from ${tableCopyWorkers} to ${numTables} to work around a Node/Postgres connection bug`
      );
    }
    const numWorkers = platform() === "win32" ? numTables : Math.min(tableCopyWorkers, numTables);
    const copyPool = pgClient(lc, upstreamURI, {
      max: numWorkers,
      connection: { ["application_name"]: "initial-sync-copy-worker" }
    });
    const copiers = startTableCopyWorkers(
      lc,
      copyPool,
      snapshot,
      numWorkers,
      numTables
    );
    try {
      createLiteTables(tx, tables, initialVersion);
      statusPublisher.publish(
        lc,
        "Initializing",
        `Copying ${numTables} upstream tables at version ${initialVersion}`,
        5e3
      );
      void copyProfiler?.start();
      const rowCounts = await Promise.all(
        tables.map(
          (table) => copiers.processReadTask(
            (db, lc2) => copy(lc2, table, copyPool, db, tx)
          )
        )
      );
      void copyProfiler?.stopAndDispose(lc, "initial-copy");
      const total = rowCounts.reduce(
        (acc, curr) => ({
          rows: acc.rows + curr.rows,
          flushTime: acc.flushTime + curr.flushTime
        }),
        { rows: 0, flushTime: 0 }
      );
      statusPublisher.publish(
        lc,
        "Indexing",
        `Creating ${indexes.length} indexes`,
        5e3
      );
      const indexStart = performance.now();
      createLiteIndices(tx, indexes);
      const index = performance.now() - indexStart;
      lc.info?.(`Created indexes (${index.toFixed(3)} ms)`);
      await addReplica(
        sql,
        shard,
        slotName,
        initialVersion,
        published,
        context
      );
      const elapsed = performance.now() - start;
      lc.info?.(
        `Synced ${total.rows.toLocaleString()} rows of ${numTables} tables in ${publications} up to ${lsn} (flush: ${total.flushTime.toFixed(3)}, index: ${index.toFixed(3)}, total: ${elapsed.toFixed(3)} ms)`
      );
    } finally {
      copiers.setDone();
      if (platform() === "win32") {
        void copyPool.end().catch((e) => lc.warn?.(`Error closing copyPool`, e));
      } else {
        await copyPool.end();
      }
    }
  } catch (e) {
    lc.warn?.(`dropping replication slot ${slotName}`, e);
    await sql`
      SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots
        WHERE slot_name = ${slotName};
    `.catch((e2) => lc.warn?.(`Unable to drop replication slot ${slotName}`, e2));
    await statusPublisher.publishAndThrowError(lc, "Initializing", e);
  } finally {
    statusPublisher.stop();
    await replicationSession.end();
    await sql.end();
  }
}
async function checkUpstreamConfig(sql) {
  const { walLevel, version } = (await sql`
      SELECT current_setting('wal_level') as "walLevel", 
             current_setting('server_version_num') as "version";
  `)[0];
  if (walLevel !== "logical") {
    throw new Error(
      `Postgres must be configured with "wal_level = logical" (currently: "${walLevel})`
    );
  }
  if (version < 15e4) {
    throw new Error(
      `Must be running Postgres 15 or higher (currently: "${version}")`
    );
  }
}
async function ensurePublishedTables(lc, sql, shard, validate = true) {
  const { database, host } = sql.options;
  lc.info?.(`Ensuring upstream PUBLICATION on ${database}@${host}`);
  await ensureShardSchema(lc, sql, shard);
  const { publications } = await getInternalShardConfig(sql, shard);
  if (validate) {
    let valid = false;
    const nonInternalPublications = publications.filter(
      (p) => !p.startsWith("_")
    );
    const exists = await sql`
      SELECT pubname FROM pg_publication WHERE pubname IN ${sql(publications)}
      `.values();
    if (exists.length !== publications.length) {
      lc.warn?.(
        `some configured publications [${publications}] are missing: [${exists.flat()}]. resyncing`
      );
    } else if (!equals(new Set(shard.publications), new Set(nonInternalPublications))) {
      lc.warn?.(
        `requested publications [${shard.publications}] differ from previouspublications [${nonInternalPublications}]. resyncing`
      );
    } else {
      valid = true;
    }
    if (!valid) {
      await sql.unsafe(dropShard(shard.appID, shard.shardNum));
      return ensurePublishedTables(lc, sql, shard, false);
    }
  }
  return { publications };
}
function startTableCopyWorkers(lc, db, snapshot, numWorkers, numTables) {
  const { init } = importSnapshot(snapshot);
  const tableCopiers = new TransactionPool(
    lc,
    READONLY,
    init,
    void 0,
    numWorkers
  );
  tableCopiers.run(db);
  lc.info?.(`Started ${numWorkers} workers to copy ${numTables} tables`);
  if (parseInt(process.versions.node) < 22) {
    lc.warn?.(
      `


Older versions of Node have a bug that results in an unresponsive
Postgres connection after running certain combinations of COPY commands.
If initial sync hangs, run zero-cache with Node v22+. This has the additional
benefit of being consistent with the Node version run in the production container image.


`
    );
  }
  return tableCopiers;
}
async function createReplicationSlot(lc, session, slotName) {
  const slot = (await session.unsafe(
    /*sql*/
    `CREATE_REPLICATION_SLOT "${slotName}" LOGICAL pgoutput`
  ))[0];
  lc.info?.(`Created replication slot ${slotName}`, slot);
  return slot;
}
function createLiteTables(tx, tables, initialVersion) {
  const columnMetadata = must(ColumnMetadataStore.getInstance(tx));
  for (const t of tables) {
    tx.exec(createLiteTableStatement(mapPostgresToLite(t, initialVersion)));
    const tableName = liteTableName(t);
    for (const [colName, colSpec] of Object.entries(t.columns)) {
      columnMetadata.insert(tableName, colName, colSpec);
    }
  }
}
function createLiteIndices(tx, indices) {
  for (const index of indices) {
    tx.exec(createLiteIndexStatement(mapPostgresToLiteIndex(index)));
  }
}
const INSERT_BATCH_SIZE = 50;
const MB = 1024 * 1024;
const MAX_BUFFERED_ROWS = 1e4;
const BUFFERED_SIZE_THRESHOLD = 8 * MB;
function makeSelectPublishedStmt(table, columns) {
  const filterConditions = Object.values(table.publications).map(({ rowFilter }) => rowFilter).filter((f) => !!f);
  return (
    /*sql*/
    `
    SELECT ${columns.map(id).join(",")} FROM ${id(table.schema)}.${id(table.name)}` + (filterConditions.length === 0 ? "" : (
      /*sql*/
      ` WHERE ${filterConditions.join(" OR ")}`
    ))
  );
}
async function copy(lc, table, dbClient, from, to) {
  const start = performance.now();
  let rows = 0;
  let flushTime = 0;
  const tableName = liteTableName(table);
  const orderedColumns = Object.entries(table.columns);
  const columnNames = orderedColumns.map(([c]) => c);
  const columnSpecs = orderedColumns.map(([_name, spec]) => spec);
  const insertColumnList = columnNames.map((c) => id(c)).join(",");
  const valuesSql = columnNames.length > 0 ? `(${"?,".repeat(columnNames.length - 1)}?)` : "()";
  const insertSql = (
    /*sql*/
    `
    INSERT INTO "${tableName}" (${insertColumnList}) VALUES ${valuesSql}`
  );
  const insertStmt = to.prepare(insertSql);
  const insertBatchStmt = to.prepare(
    insertSql + `,${valuesSql}`.repeat(INSERT_BATCH_SIZE - 1)
  );
  const selectStmt = makeSelectPublishedStmt(table, columnNames);
  const valuesPerRow = columnSpecs.length;
  const valuesPerBatch = valuesPerRow * INSERT_BATCH_SIZE;
  const pendingValues = Array.from({
    length: MAX_BUFFERED_ROWS * valuesPerRow
  });
  let pendingRows = 0;
  let pendingSize = 0;
  function flush() {
    const start2 = performance.now();
    const flushedRows = pendingRows;
    const flushedSize = pendingSize;
    let l = 0;
    for (; pendingRows > INSERT_BATCH_SIZE; pendingRows -= INSERT_BATCH_SIZE) {
      insertBatchStmt.run(pendingValues.slice(l, l += valuesPerBatch));
    }
    for (; pendingRows > 0; pendingRows--) {
      insertStmt.run(pendingValues.slice(l, l += valuesPerRow));
    }
    for (let i = 0; i < flushedRows; i++) {
      pendingValues[i] = void 0;
    }
    pendingSize = 0;
    rows += flushedRows;
    const elapsed2 = performance.now() - start2;
    flushTime += elapsed2;
    lc.debug?.(
      `flushed ${flushedRows} ${tableName} rows (${flushedSize} bytes) in ${elapsed2.toFixed(3)} ms`
    );
  }
  lc.info?.(`Starting copy stream of ${tableName}:`, selectStmt);
  const pgParsers = await getTypeParsers(dbClient, { returnJsonAsString: true });
  const parsers = columnSpecs.map((c) => {
    const pgParse = pgParsers.getTypeParser(c.typeOID);
    return (val) => liteValue(
      pgParse(val),
      c.dataType,
      JSON_STRINGIFIED
    );
  });
  const tsvParser = new TsvParser();
  let col = 0;
  await pipeline(
    await from.unsafe(`COPY (${selectStmt}) TO STDOUT`).readable(),
    new Writable({
      highWaterMark: BUFFERED_SIZE_THRESHOLD,
      write(chunk, _encoding, callback) {
        try {
          for (const text of tsvParser.parse(chunk)) {
            pendingSize += text === null ? 4 : text.length;
            pendingValues[pendingRows * valuesPerRow + col] = text === null ? null : parsers[col](text);
            if (++col === parsers.length) {
              col = 0;
              if (++pendingRows >= MAX_BUFFERED_ROWS - valuesPerRow || pendingSize >= BUFFERED_SIZE_THRESHOLD) {
                flush();
              }
            }
          }
          callback();
        } catch (e) {
          callback(e instanceof Error ? e : new Error(String(e)));
        }
      },
      final: (callback) => {
        try {
          flush();
          callback();
        } catch (e) {
          callback(e instanceof Error ? e : new Error(String(e)));
        }
      }
    })
  );
  const elapsed = performance.now() - start;
  lc.info?.(
    `Finished copying ${rows} rows into ${tableName} (flush: ${flushTime.toFixed(3)} ms) (total: ${elapsed.toFixed(3)} ms) `
  );
  return { rows, flushTime };
}
export {
  INSERT_BATCH_SIZE,
  createReplicationSlot,
  initialSync,
  makeSelectPublishedStmt
};
//# sourceMappingURL=initial-sync.js.map
