import { readFile, readdir } from 'node:fs/promises';
import path from 'node:path';
import pg from 'pg';

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,
      applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
    )
  `);

  if (direction === 'up') {
    const appliedResult = await client.query('SELECT version FROM schema_migrations');
    const applied = new Set(appliedResult.rows.map((row) => String(row.version)));
    for (const file of files) {
      const version = file.replace(/\.up\.sql$/, '');
      if (applied.has(version)) continue;
      const sql = await readFile(path.join(migrationsDirectory, file), 'utf8');
      await applyTransaction(async () => {
        await client.query(sql);
        await client.query('INSERT INTO schema_migrations (version) VALUES ($1)', [version]);
      });
      console.log(`applied ${file}`);
    }
  } else {
    const requestedVersion = process.argv[3];
    const latestResult = await client.query(
      'SELECT version 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 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;
  }
}