import partition from "lodash/partition";
import sum from "lodash/sum";
import { getPhysicalReplicas } from "./getPhysicalReplicas";
import { getRowCount } from "./getRowCount";
import type { TableInfo } from "./getTablesInSchema";
import { log, progress } from "./logging";
import { subName } from "./names";
import { pluralize } from "./pluralize";
import { pollDelay } from "./pollDelay";
import { sql } from "./quote";
import { runSql } from "./runSql";
import { secToHuman } from "./secToHuman";
import { tookToHuman } from "./tookToHuman";

const MAX_INITIALIZING_TABLES = 4;

const STATES: Record<string, string | unknown> = {
  i: "initializing",
  d: "data is being copied",
  f: "finished table copy",
  s: "synchronized",
  r: "ready",
};

/**
 * Waits until backfill of the destination tables finishes.
 * https://www.postgresql.org/docs/14/catalog-pg-subscription-rel.html
 */
export async function waitUntilBackfillCompletes(
  {
    fromDsn,
    tables,
    toDsn,
    schema,
  }: {
    fromDsn: string;
    tables: TableInfo[];
    toDsn: string;
    schema: string;
  },
  throwIfAborted: () => void,
): Promise<void> {
  if (tables.length === 0) {
    return;
  }

  const waitStartedAt = performance.now();
  const tableRowCountsTo = new Map<string, number>();
  const firstSeenActiveTables = new Map<string, number>();
  const stats: string[] = [];

  while (true) {
    stats.splice(0, stats.length);

    const backfillingTables = await runSql(
      toDsn,
      sql`
        SELECT relname, srsubstate
        FROM pg_subscription_rel
        JOIN pg_subscription ON pg_subscription.oid = srsubid
        JOIN pg_class ON pg_class.oid = srrelid
        JOIN pg_database ON pg_database.oid = subdbid AND datname = current_database()
        WHERE subname = ${subName(schema)} AND srsubstate <> 'r'
      `,
    );

    for (const [table] of backfillingTables) {
      if (!tableRowCountsTo.has(table)) {
        tableRowCountsTo.set(table, 0);
      }
    }

    const [activeTables, initializingTables] = partition(
      backfillingTables,
      ([, state]) => state !== "i",
    );

    for (const [table, state] of activeTables) {
      if (!firstSeenActiveTables.has(table)) {
        firstSeenActiveTables.set(table, performance.now());
      }

      const countFrom =
        (await getRowCount({
          dsn: fromDsn,
          schema,
          table,
          method: "explain",
        })) ?? 0;
      const countTo =
        (await getRowCount({
          dsn: toDsn,
          schema,
          table,
          method: "pg_stat_progress_copy",
        })) ??
        tableRowCountsTo.get(table) ??
        0;
      tableRowCountsTo.set(table, countTo);
      const percent =
        countFrom === 0
          ? "100"
          : Math.min(100, Math.round((countTo / countFrom) * 100));
      const stateName = STATES[state] ?? state;
      const approx = countTo > countFrom ? "~" : "";
      const elapsedSec = Math.round(
        (performance.now() - firstSeenActiveTables.get(table)!) / 1000,
      );
      const leftSec = Math.round(
        Math.max(
          countFrom === 0 || countTo === 0
            ? 0
            : (elapsedSec / countTo) * (countFrom - countTo),
          0,
        ),
      );
      stats.push(
        `${table}: ${stateName}, ${approx}${countTo} of ${approx}${countFrom}, ${approx}${percent}%: elapsed=${secToHuman(elapsedSec)} left=${approx}${secToHuman(leftSec)}`,
      );
      throwIfAborted();
    }

    if (initializingTables.length > 0) {
      stats.push(
        initializingTables
          .map(([table]) => table)
          .slice(0, MAX_INITIALIZING_TABLES)
          .join(", ") +
          (initializingTables.length > MAX_INITIALIZING_TABLES
            ? ` and ${pluralize(initializingTables.length - MAX_INITIALIZING_TABLES, "other table")}`
            : "") +
          ": waiting in max_sync_workers_per_subscription queue",
      );
    }

    const physicalReplicasFrom = await getPhysicalReplicas({ dsn: fromDsn });
    const physicalReplicasTo = await getPhysicalReplicas({ dsn: toDsn });

    for (const [
      srcDst,
      {
        clientAddr,
        userName,
        applicationName,
        lagSec,
        feedbackAgoSec,
        masterLsn,
        replicaLsn,
        gap,
      },
    ] of [
      ...physicalReplicasFrom.map((replica) => ["src", replica] as const),
      ...physicalReplicasTo.map((replica) => ["dst", replica] as const),
    ] as const) {
      const prefix = `physical replica of the ${srcDst} ${userName}@${clientAddr} (${applicationName})`;
      stats.push(
        `${prefix}: lag=${secToHuman(lagSec)} master=${masterLsn} replica=${replicaLsn} gap=${gap} feedbackAgo=${secToHuman(feedbackAgoSec)}`,
      );
    }

    if (backfillingTables.length === 0) {
      // Keep the data in stats.
      break;
    }

    progress(
      "The logical tablesync worker is still backfilling:",
      ...stats.map((str) => `- ${str}`),
    );
    await pollDelay();
    throwIfAborted();
  }

  progress.clear();
  const totalTuples = sum([...tableRowCountsTo.values()]);
  log.shellCmd({
    comment: `Logical tablesync worker backfill completed in ${tookToHuman(waitStartedAt)}`,
    cmd: "",
    input: [
      `backfilled ${pluralize(tableRowCountsTo.size, "table")}, ~${pluralize(totalTuples, "tuple")}`,
      ...stats,
    ]
      .map((str) => `+ ${str}`)
      .join("\n"),
  });
}