import { readFile, readdir } from 'node:fs/promises';
import path from 'node:path';
import pg from 'pg';
import {
  assertAppliedMigrationsPresent,
  assertMigrationChecksum,
  migrationChecksum,
} from './migration-checksum.mjs';

const direction = process.argv[2];
if (direction !== 'up' && direction !== 'down') {
  throw new Error('Usage: node scripts/migrate.mjs <up|down>');
}

const connectionString = process.env.DATABASE_MIGRATION_URL ?? process.env.DATABASE_URL;
if (!connectionString) throw new Error('DATABASE_MIGRATION_URL or DATABASE_URL is required');
if (direction === 'down' && process.env.ALLOW_MIGRATION_DOWN !== 'true') {
  throw new Error(
    'Refusing destructive rollback. Set ALLOW_MIGRATION_DOWN=true to roll back the latest migration.',
  );
}

const migrationsDirectory = path.resolve('migrations');
const files = (await readdir(migrationsDirectory))
  .filter((name) => name.endsWith('.up.sql'))
  .sort();
const client = new pg.Client({ connectionString, application_name: 'krona-staff-migrator' });
await client.connect();

const advisoryLockId = 7_347_821_009;
try {
  await client.query('SELECT pg_advisory_lock($1)', [advisoryLockId]);
  await client.query(`
    CREATE TABLE IF NOT EXISTS schema_migrations (
      version TEXT PRIMARY KEY,
      checksum TEXT,
      applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
    )
  `);
  await client.query('ALTER TABLE schema_migrations ADD COLUMN IF NOT EXISTS checksum TEXT');

  if (direction === 'up') {
    const appliedResult = await client.query('SELECT version, checksum FROM schema_migrations');
    const applied = new Map(
      appliedResult.rows.map((row) => [String(row.version), row.checksum && String(row.checksum)]),
    );
    const availableVersions = new Set(files.map((file) => file.replace(/\.up\.sql$/, '')));
    assertAppliedMigrationsPresent(applied.keys(), availableVersions);
    for (const file of files) {
      const version = file.replace(/\.up\.sql$/, '');
      const content = await readFile(path.join(migrationsDirectory, file));
      const checksum = migrationChecksum(content);
      if (applied.has(version)) {
        const storedChecksum = applied.get(version);
        if (storedChecksum) {
          assertMigrationChecksum(version, storedChecksum, checksum);
        } else {
          await client.query('UPDATE schema_migrations SET checksum = $2 WHERE version = $1', [
            version,
            checksum,
          ]);
          console.warn(`baselined checksum for legacy migration ${version}`);
        }
        continue;
      }
      const sql = content.toString('utf8');
      await applyTransaction(async () => {
        await client.query(sql);
        await client.query('INSERT INTO schema_migrations (version, checksum) VALUES ($1, $2)', [
          version,
          checksum,
        ]);
      });
      console.log(`applied ${file}`);
    }
  } else {
    const requestedVersion = process.argv[3];
    const latestResult = await client.query(
      'SELECT version, checksum FROM schema_migrations ORDER BY version DESC LIMIT 1',
    );
    const latestVersion = latestResult.rows[0]?.version;
    if (!latestVersion) {
      console.log('no migration to roll back');
    } else {
      if (requestedVersion && requestedVersion !== latestVersion) {
        throw new Error(
          `Only the latest migration can be rolled back. Latest is ${latestVersion}.`,
        );
      }
      const downFile = `${latestVersion}.down.sql`;
      const upContent = await readFile(path.join(migrationsDirectory, `${latestVersion}.up.sql`));
      const currentChecksum = migrationChecksum(upContent);
      if (!latestResult.rows[0].checksum) {
        throw new Error(
          `Migration ${latestVersion} has no checksum. Run the up command once to establish the legacy baseline before rollback.`,
        );
      }
      assertMigrationChecksum(
        latestVersion,
        String(latestResult.rows[0].checksum),
        currentChecksum,
      );
      const sql = await readFile(path.join(migrationsDirectory, downFile), 'utf8');
      await applyTransaction(async () => {
        await client.query(sql);
        await client.query('DELETE FROM schema_migrations WHERE version = $1', [latestVersion]);
      });
      console.log(`rolled back ${downFile}`);
    }
  }
} finally {
  await client.query('SELECT pg_advisory_unlock($1)', [advisoryLockId]).catch(() => undefined);
  await client.end();
}

async function applyTransaction(run) {
  await client.query('BEGIN');
  try {
    await run();
    await client.query('COMMIT');
  } catch (error) {
    try {
      await client.query('ROLLBACK');
    } catch (rollbackError) {
      throw new AggregateError([error, rollbackError], 'Migration failed and rollback also failed');
    }
    throw error;
  }
}