import { execFile as execFileCallback } from "child_process";
import { program } from "commander";
import { diffLines } from "diff";
import dotenv from "dotenv";
import { readdirSync, readFileSync, writeFileSync } from "fs";
import { dirname, join } from "path";
import pg from "pg";
import { fileURLToPath } from "url";
import { promisify } from "util";

dotenv.config();

const { Pool } = pg;
const execFile = promisify(execFileCallback);

// given a directory of sql files with names like 001_create_table.sql, 002_add_column.sql, etc.
// return a list of { id, filename, sql } in ascending order of id
const getMigrations = async (path) => {
  const files = readdirSync(path);
  const migrations = [];
  for (const filename of files) {
    if (!filename.endsWith(".sql") || filename.startsWith(".")) {
      continue;
    }
    // use a regex to extract the id from the filename so we can report errors
    const match = filename.match(/^(\d+)_.*\.sql$/);
    if (!match) {
      throw new Error(
        `Expected migration filename like {number}_{description}.sql but got ${filename}`
      );
    }
    const id = Number(match[1]);
    const sql = readFileSync(`${path}/${filename}`, "utf8");
    migrations.push({ id, filename, sql });
  }

  migrations.sort((a, b) => a.id - b.id);

  let lastId = null;
  for (let i = 0; i < migrations.length; i++) {
    if (migrations[i].id === lastId) {
      throw new Error(
        `Duplicate migration id ${lastId} in ${migrations[i].filename} and ${
          migrations[i - 1].filename
        }`
      );
    }
    lastId = migrations[i].id;
  }
  return migrations;
};

const splitMigration = (filename, sql) => {
  const sqlMatch = sql.match(/^(\s*--up\s*\n.*?)(\n\s*--down\s*\n.*)?$/s);
  if (!sqlMatch) {
    throw new Error(
      `Expected migration ${filename} to contain --up and optional --down blocks`
    );
  }
  return { up: sqlMatch[1], down: sqlMatch[2] };
};

const getPendingMigrations = async (
  connectionInTransaction,
  migrations,
  allowChanges = false
) => {
  const removed = [];
  const added = [];

  const result = await connectionInTransaction.query(
    "SELECT id, filename, content FROM migrations ORDER BY id ASC"
  );

  for (const row of result.rows) {
    const migration = migrations.find((m) => m.id === row.id);
    if (allowChanges) {
      if (
        !migration ||
        migration.filename !== row.filename ||
        migration.sql !== row.content
      ) {
        removed.push({ id: row.id, filename: row.filename, sql: row.content });
      }
    } else {
      if (!migration) {
        throw new Error(
          `Migration ${row.id} has been applied but is missing from the migrations directory`
        );
      }
      if (migration.filename !== row.filename) {
        throw new Error(
          `Migration ${row.id} has been applied but has a different filename`
        );
      }
      if (migration.sql !== row.content) {
        throw new Error(
          `Migration ${row.id} has been applied but has different content`
        );
      }
    }
  }

  const maxId =
    result.rows.length > 0 ? result.rows[result.rows.length - 1].id : -1;
  for (const migration of migrations) {
    if (result.rows.find((row) => row.id === migration.id)) {
      continue;
    }
    if (!allowChanges && migration.id <= maxId) {
      throw new Error(
        `Unapplied migration ${migration.id} has id lower than the last applied migration`
      );
    }
    added.push(migration);
  }

  return { removed, added };
};

class SqlError extends Error {}

// run sql block with nice error messages
const runSql = async (connection, sqlFilename, sql) => {
  try {
    await connection.query(sql);
  } catch (e) {
    if (e.position?.match(/\d+/)) {
      let position = Number(e.position) - 1;
      const lines = sql.split("\n");
      for (let i = 0; i < lines.length; i++) {
        const line = lines[i];
        if (position < line.length) {
          throw new SqlError(
            `Error at ${sqlFilename}:${i + 1}:\n${line}\n${" ".repeat(
              position
            )}^\n${e.message}`
          );
        }
        position -= line.length + 1;
      }
    }
    throw e;
  }
};

const getPostgresPool = async (dbName?) => {
  const config = {
    user: process.env.PGUSER,
    host: process.env.PGHOST,
    database: dbName || process.env.PGDATABASE,
    password: process.env.PGPASSWORD,
    port: Number(process.env.PGPORT || 5432),
    max: 1,
  } as any;
  if (process.env.PGSSLMODE === "require") {
    config.ssl = { rejectUnauthorized: false };
  }

  const pool = new Pool(config);

  return pool;
};

const withPostgresPool = async (dbName, fn) => {
  const pool = await getPostgresPool(dbName);
  try {
    return await fn(pool);
  } finally {
    await pool.end();
  }
};

const withPostgresConnection = async (dbName, fn) => {
  return await withPostgresPool(dbName, async (pool) => {
    const conn = await pool.connect();
    try {
      return await fn(conn);
    } finally {
      await conn.release();
    }
  });
};

const withLockedMigrationsTable = async (dbName, fn) => {
  return await withPostgresConnection(dbName, async (conn) => {
    await conn.query("BEGIN");
    await conn.query(`CREATE TABLE IF NOT EXISTS migrations (
      id integer PRIMARY KEY,
      filename text NOT NULL,
      "timestamp" timestamp without time zone DEFAULT now() NOT NULL,
      content text NOT NULL
    )`);
    await conn.query("LOCK TABLE migrations IN EXCLUSIVE MODE");
    return await fn(conn);
  });
};

// apply all pending migrations, allowing only new migrations to be added
const migrate = async (path, shouldRunMigration, dbName?) => {
  const migrations = await getMigrations(path);
  return await withLockedMigrationsTable(dbName, async (conn) => {
    const { added: pendingMigrations } = await getPendingMigrations(
      conn,
      migrations
    );
    if (pendingMigrations.length === 0) {
      console.log(
        `No new migrations to ${shouldRunMigration ? "apply" : "adopt"}`
      );
    }
    for (const migration of pendingMigrations) {
      console.log(
        `${shouldRunMigration ? "Applying" : "Adopting"} migration ${
          migration.filename
        }`
      );
      const { up } = splitMigration(migration.filename, migration.sql);
      if (shouldRunMigration) {
        await runSql(conn, migration.filename, up);
      }
      await conn.query(
        "INSERT INTO migrations (id, filename, content) VALUES ($1, $2, $3)",
        [migration.id, migration.filename, migration.sql]
      );
    }
    await conn.query("COMMIT");
  });
};

// apply migrations in such a way that the db migration history matches the migrations directory
const sync = async (migrationsPath, schemaPath?, dbName?) => {
  const migrations = await getMigrations(migrationsPath);
  const migratedSchema = await checkMigrations(migrationsPath);
  if (!migratedSchema) {
    process.exitCode = 1;
    return;
  }
  await withLockedMigrationsTable(dbName, async (conn) => {
    const { added, removed } = await getPendingMigrations(
      conn,
      migrations,
      true
    );
    if (added.length === 0 && removed.length === 0) {
      console.log("No migrations to sync");
      return;
    }
    // find the migration added or removed with the lowest id and roll back to just before that id
    const rollbackId = Math.min(...added.concat(removed).map((m) => m.id));
    const result = await conn.query(
      "SELECT id, filename, content FROM migrations WHERE id >= $1 ORDER BY id DESC",
      [rollbackId]
    );
    const migrationsToRollBack = result.rows;
    for (const migration of migrationsToRollBack) {
      console.log(`Rolling back migration ${migration.filename}`);
      const { down } = splitMigration(migration.filename, migration.content);
      if (down !== undefined) {
        await runSql(conn, `(migrations table)${migration.filename}`, down);
      }
      await conn.query("DELETE FROM migrations WHERE id = $1", [migration.id]);
    }
    const { added: addedAfterRollBack } = await getPendingMigrations(
      conn,
      migrations
    );
    for (const migration of addedAfterRollBack) {
      console.log(`Applying migration ${migration.filename}`);
      const { up } = splitMigration(migration.filename, migration.sql);
      await runSql(conn, migration.filename, up);
      await conn.query(
        "INSERT INTO migrations (id, filename, content) VALUES ($1, $2, $3)",
        [migration.id, migration.filename, migration.sql]
      );
    }
    await conn.query("COMMIT");
  });
  if (schemaPath) {
    writeFileSync(schemaPath, migratedSchema);
  }
};

const cloneProductionDb = async (
  productionDbHost,
  productionDbPort,
  productionDbPassword,
  migrationsPath,
  cloneScriptPath,
  dbName?
) => {
  const cloneScript = readFileSync(cloneScriptPath, "utf8");
  await sync(migrationsPath, null, dbName);
  if (process.exitCode === 1) {
    return;
  }
  console.log("Running production clone script (this will take a minute)...");
  await withPostgresConnection(dbName, async (conn) => {
    await conn.query(`SET wavtool.production_db_host TO '${productionDbHost}'`);
    await conn.query(`SET wavtool.production_db_port TO '${productionDbPort}'`);
    await conn.query(
      `SET wavtool.production_db_password TO '${productionDbPassword}'`
    );
    conn.on("notice", msg => {
      if(msg?.routine !== 'exec_stmt_raise') {
        return;
      }
      console.log(msg.message);
    });
    await runSql(conn, cloneScriptPath, cloneScript);
  });
};

// use pg_dump to dump database schema
const dumpSchema = async (db?) => {
  const pgDumpCommand = process.env.PG_DUMP_COMMAND || "pg_dump";
  const [pgDumpExecutable, ...pgDumpArgs] = pgDumpCommand.split(" ");
  const { stdout, stderr } = await execFile(pgDumpExecutable, [
    ...pgDumpArgs,
    "--no-acl",
    "--no-owner",
    "--schema-only",
    "--no-comments",
    ...(db ? [db] : []),
  ]);
  if (stderr) {
    console.error(stderr);
  }
  return stdout.replace(/^-- Dumped from.*\n-- Dumped by.*\n/m, "");
};

// use pg_dump to determine diff between two databases
const diffSchemas = async (dbA, dbB) => {
  const schemaA = await dumpSchema(dbA);
  const schemaB = await dumpSchema(dbB);
  return diffLines(schemaA, schemaB, { ignoreWhitespace: true }).filter(
    (d) => d.added || d.removed
  );
};

// generate a temporary database name
const generateTempDbName = () => {
  return `temp_${Math.random().toString(36).substring(2, 15)}`;
};

const tryDropDb = async (conn, dbName) => {
  try {
    await conn.query(`DROP DATABASE IF EXISTS ${dbName}`);
  } catch (e) {
    console.log(`Failed to drop database ${dbName}: ${e.message}`);
  }
};

const compareDatabaseToSchema = async (schemaPath, dbName?) => {
  const schema = readFileSync(schemaPath, "utf8");
  const schemaDb = generateTempDbName();
  return await withPostgresConnection(null, async (conn) => {
    try {
      await conn.query(`CREATE DATABASE ${schemaDb}`);
      await withPostgresConnection(schemaDb, async (schemaConn) => {
        await runSql(schemaConn, schemaPath, schema);
      });

      const diff = await diffSchemas(dbName, schemaDb);
      if (diff.length > 0) {
        console.error(
          `Database and schema do not match. Diff:\n${diff
            .map((d) => (d.added ? "+ " : "- ") + d.value)
            .join("")}`
        );
        process.exitCode = 1;
      }
    } finally {
      await tryDropDb(conn, schemaDb);
    }
  });
};

const checkMigrations = async (migrationsPath) => {
  const migrationsDb = generateTempDbName();
  return await withPostgresConnection(null, async (conn) => {
    try {
      await conn.query(`CREATE DATABASE ${migrationsDb}`);

      const migrations = await getMigrations(migrationsPath);
      const allCorrectlyInverse = await withLockedMigrationsTable(
        migrationsDb,
        async (conn) => {
          await conn.query("COMMIT");

          const { added: pendingMigrations } = await getPendingMigrations(
            conn,
            migrations
          );

          for (const migration of pendingMigrations) {
            const { up, down } = splitMigration(
              migration.filename,
              migration.sql
            );

            const schemaBefore = down && (await dumpSchema(migrationsDb));

            await conn.query("BEGIN");
            await runSql(conn, migration.filename, up);
            await conn.query("COMMIT");

            if (!down) {
              continue;
            }

            await conn.query("BEGIN");
            await runSql(conn, migration.filename, down);
            await conn.query("COMMIT");

            const schemaAfter = await dumpSchema(migrationsDb);

            const diff = diffLines(schemaBefore, schemaAfter, {
              ignoreWhitespace: true,
            }).filter((d) => d.added || d.removed);
            if (diff.length > 0) {
              console.error(
                `Migration ${
                  migration.filename
                } --down section does not invert --up section. Diff:\n${diff
                  .map((d) => (d.added ? "+ " : "- ") + d.value)
                  .join("")}`
              );
              return false;
            }

            await conn.query("BEGIN");
            await runSql(conn, migration.filename, up);
            await conn.query("COMMIT");
          }

          return true;
        }
      );

      if (!allCorrectlyInverse) {
        return false;
      }

      return await dumpSchema(migrationsDb);
    } finally {
      await tryDropDb(conn, migrationsDb);
    }
  });
};

const compareSchemaToMigrations = async (schemaPath, migrationsPath) => {
  const schema = readFileSync(schemaPath, "utf8");
  const schemaDb = generateTempDbName();

  const migratedSchema = await checkMigrations(migrationsPath);
  if (!migratedSchema) {
    process.exitCode = 1;
    return;
  }
  return await withPostgresConnection(null, async (conn) => {
    try {
      await conn.query(`CREATE DATABASE ${schemaDb}`);
      await withPostgresConnection(schemaDb, async (schemaConn) => {
        await runSql(schemaConn, schemaPath, schema);
      });

      const normalizedSchema = await dumpSchema(schemaDb);
      const diff = diffLines(normalizedSchema, migratedSchema, {
        ignoreWhitespace: true,
      }).filter((d) => d.added || d.removed);

      if (diff.length > 0) {
        console.error(
          `Schema and migrations do not match. Diff:\n${diff
            .map((d) => (d.added ? "+ " : "- ") + d.value)
            .join("")}`
        );
        console.error(
          `\nHINT: If you're seeing this error in CI, be sure to run\n` +
            `"yarn migrate sync" locally and commit the resulting changes\n` +
            `to db/wavtool_schema.sql\n`
        );
        process.exitCode = 1;
      }
    } finally {
      await tryDropDb(conn, schemaDb);
    }
  });
};

const withSqlErrorDisplay = async (fn) => {
  try {
    await fn();
  } catch (e) {
    if (e instanceof SqlError) {
      console.error(e.message);
      process.exitCode = 1;
    } else {
      throw e;
    }
  }
};

const __dirname = dirname(fileURLToPath(import.meta.url));
const defaultMigrationsPath = join(__dirname, "..", "..", "migrations");
const defaultSchemaPath = join(__dirname, "..", "..", "wavtool_schema.sql");

// Prevent accidents
if (
  (process.env.PGHOST || "").includes("us-east-1.rds.amazonaws.com") &&
  !process.env.REALLY_RUN_AGAINST_PRODUCTION
) {
  console.log(
    "You are about to run migrations against the production database."
  );
  console.log(
    "If you really want to do this, set REALLY_RUN_AGAINST_PRODUCTION=1 in your environment."
  );
  process.exit(1);
}

program.name("migrate");

program
  .command("apply")
  .description("Apply all new migrations")
  .option(
    "--migrations-path <migrationsPath>",
    "Path to migrations directory",
    defaultMigrationsPath
  )
  .action(async (opts) => {
    await withSqlErrorDisplay(async () => {
      await migrate(opts.migrationsPath, true);
    });
  });

program
  .command("adopt")
  .description("Record all new migrations as applied without running them")
  .option(
    "--migrations-path <migrationsPath>",
    "Path to migrations directory",
    defaultMigrationsPath
  )
  .action(async (opts) => {
    await withSqlErrorDisplay(async () => {
      await migrate(opts.migrationsPath, false);
    });
  });

program
  .command("sync")
  .description("Sync database state to migrations")
  .option(
    "--migrations-path <migrationsPath>",
    "Path to migrations directory",
    defaultMigrationsPath
  )
  .option(
    "--schema-path <schemaPath>",
    "Path to full schema file (will be updated)",
    defaultSchemaPath
  )
  .action(async (opts) => {
    await withSqlErrorDisplay(async () => {
      await sync(opts.migrationsPath, opts.schemaPath);
    });
  });

program
  .command("compare")
  .description("Compare full schema to migrations")
  .option(
    "--migrations-path <migrationsPath>",
    "Path to migrations directory",
    defaultMigrationsPath
  )
  .option(
    "--schema-path <schemaPath>",
    "Path to full schema file",
    defaultSchemaPath
  )
  .action(async (opts) => {
    await withSqlErrorDisplay(async () => {
      await compareSchemaToMigrations(opts.schemaPath, opts.migrationsPath);
    });
  });

program
  .command("comparedb")
  .description("Compare database to schema")
  .option(
    "--schema-path <schemaPath>",
    "Path to full schema file",
    defaultSchemaPath
  )
  .action(async (opts) => {
    await withSqlErrorDisplay(async () => {
      await compareDatabaseToSchema(opts.schemaPath);
    });
  });

program
  .command("create")
  .description("Create new migration")
  .option(
    "--migrations-path <migrationsPath>",
    "Path to migrations directory",
    defaultMigrationsPath
  )
  .argument("<name>", "Name of migration")
  .action(async (name, opts) => {
    const path = join(
      opts.migrationsPath,
      `${Math.round(Date.now() / 1000)}_${name}.sql`
    );
    writeFileSync(path, "--up\n\n--down\n\n");
  });

program
  .command("clone-prod")
  .description(
    "Clone production database. Environment must supply PRODUCTION_PGHOST, PRODUCTION_PGPORT, and PRODUCTION_PGPASSWORD."
  )
  .option(
    "--migrations-path <migrationsPath>",
    "Path to migrations directory",
    defaultMigrationsPath
  )
  .option(
    "--clone-script-path <cloneScriptPath>",
    "Path to production clone script",
    join(__dirname, "..", "..", "clone_prod.sql")
  )
  .action(async (opts) => {
    const productionDbHost = process.env.PRODUCTION_PGHOST;
    const productionDbPort = process.env.PRODUCTION_PGPORT || 5432;
    const productionDbPassword = process.env.PRODUCTION_PGPASSWORD;
    if (!productionDbHost || !productionDbPassword) {
      console.error(
        "PRODUCTION_PGHOST and PRODUCTION_PGPASSWORD must be set in environment"
      );
      process.exit(1);
    }
    await withSqlErrorDisplay(async () => {
      await cloneProductionDb(
        productionDbHost,
        productionDbPort,
        productionDbPassword,
        opts.migrationsPath,
        opts.cloneScriptPath
      );
    });
  });

program.parse();