#!/usr/bin/env npx tsx
/**
 * T004b Migration Script — In-place PII Encryption of Existing Rows
 * Brief: docs/regulatory/T004B_DESIGN_BRIEF.md v1.2-RATIFIED
 *
 * Usage:
 *   DRY_RUN=true  npx tsx scripts/t004b-migrate.ts   (default; writes snapshot, NO DB writes)
 *   DRY_RUN=false npx tsx scripts/t004b-migrate.ts   (live; requires operator authorisation)
 *
 * SECURITY: NO plaintext PII values appear in ANY log output — not on normal path,
 * dry-run path, gate-failure path, or stale-overwrite path.
 * All abort messages reference the snapshot file path only.
 * The encrypted .enc snapshot file is the sole record of original plaintext values.
 */

import pg from "pg";
import * as fs from "fs";
import {
  encryptPii,
  decryptPii,
  isEncrypted,
  computeLookupHash,
  type PiiContext,
} from "../server/lib/pii-encryption";

// ─── Configuration ────────────────────────────────────────────────────────────

const DRY_RUN = process.env.DRY_RUN !== "false"; // default true — safe mode
const TIMESTAMP = Date.now();
const SNAPSHOT_PATH = `./t004b-snapshot-${TIMESTAMP}.enc`;
const PLATFORM_CTX: PiiContext = { type: "platform" };

// ─── Table config ─────────────────────────────────────────────────────────────

interface Phase1TableConfig {
  table: string;
  primaryCol: string; // isEncrypted() idempotency sentinel (first listed PII col)
  piiCols: string[];  // ALL PII columns — encrypted + included in WHERE predicate
  contextFn: () => PiiContext;
}

const PHASE1_TABLES: Phase1TableConfig[] = [
  {
    table: "suspicious_activity_reports",
    primaryCol: "subject_name",
    piiCols: ["subject_name", "subject_id", "description"],
    contextFn: () => PLATFORM_CTX,
  },
  {
    table: "explanation_dossiers",
    primaryCol: "subject_name",
    piiCols: ["subject_name", "subject_id", "contact_information"],
    contextFn: () => PLATFORM_CTX,
  },
  {
    table: "vendor_risk_register",
    primaryCol: "contact_email",
    piiCols: ["contact_email"],
    contextFn: () => PLATFORM_CTX,
  },
  {
    table: "key_inventory",
    primaryCol: "owner_email",
    piiCols: ["owner_email"],
    contextFn: () => PLATFORM_CTX,
  },
  {
    table: "model_cards_v2",
    primaryCol: "owner_email",
    piiCols: ["owner_email"],
    contextFn: () => PLATFORM_CTX,
  },
  {
    table: "users",
    primaryCol: "full_name",
    piiCols: ["full_name"],
    contextFn: () => PLATFORM_CTX,
    // M-G: full_name is notNull at DB level but treated as nullable-safe.
    // IS NOT DISTINCT FROM in WHERE predicate handles both cases correctly.
  },
];

// ─── Snapshot types ───────────────────────────────────────────────────────────

// Phase 1 snapshot: table → rowId → { col: plaintext | null }
type Phase1Snap = Record<string, Record<string, Record<string, string | null>>>;

// Phase 2 snapshot: table → rowId → { blobCol: plaintext, tenantId: string | null }
// tenantId is NOT PII but is stored as metadata for recovery key derivation.
type Phase2Row = { blobValue: string; tenantId: string | null };
type Phase2Snap = Record<string, Record<string, Phase2Row>>;

interface FullSnapshot {
  phase1: Phase1Snap;
  phase2: Phase2Snap;
}

// ─── Stats ────────────────────────────────────────────────────────────────────

const stats = {
  migrated: 0,
  skipped: 0,
  hmacBackfilled: 0,
  nullTenantSkipped: 0, // Phase 2 rows skipped because tenant_id IS NULL (cannot derive HMAC key)
};

// ─── Abort error (never includes plaintext) ───────────────────────────────────

class MigrationAbortError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "MigrationAbortError";
  }
}

// ─── Utilities ────────────────────────────────────────────────────────────────

function buildUpdateSQL(
  table: string,
  piiCols: string[]
): { sql: string; paramBuilder: (rowId: string, snapVals: (string | null)[], cipherVals: (string | null)[]) => unknown[] } {
  // Param layout: $1=rowId, $2..$N+1=snapVals (WHERE), $N+2..$2N+1=cipherVals (SET)
  // where N = piiCols.length
  const n = piiCols.length;
  const setClauses = piiCols.map((col, i) => `${col} = $${n + i + 2}`).join(", ");
  const whereClauses = piiCols.map((col, i) => `${col} IS NOT DISTINCT FROM $${i + 2}`).join(" AND ");
  const sql = `UPDATE ${table} SET ${setClauses} WHERE id = $1 AND ${whereClauses}`;
  const paramBuilder = (rowId: string, snapVals: (string | null)[], cipherVals: (string | null)[]) =>
    [rowId, ...snapVals, ...cipherVals];
  return { sql, paramBuilder };
}

// ─── Phase 0 — Build complete snapshot ───────────────────────────────────────

async function phase0BuildSnapshot(pool: pg.Pool): Promise<FullSnapshot> {
  const phase1: Phase1Snap = {};
  const phase2: Phase2Snap = {};

  // Phase 1 tables
  for (const cfg of PHASE1_TABLES) {
    phase1[cfg.table] = {};
    const colList = ["id", ...cfg.piiCols].join(", ");
    const res = await pool.query(
      `SELECT ${colList} FROM ${cfg.table}
       WHERE ${cfg.primaryCol} NOT LIKE 'ENC:%' AND ${cfg.primaryCol} IS NOT NULL`
    );
    for (const row of res.rows) {
      const snap: Record<string, string | null> = {};
      for (const col of cfg.piiCols) {
        snap[col] = (row[col] as string | null) ?? null;
      }
      phase1[cfg.table][row.id as string] = snap;
    }
    const count = Object.keys(phase1[cfg.table]).length;
    console.log(`[t004b] [phase-0] ${cfg.table}: ${count} row(s) to migrate`);
  }

  // Phase 2 — HMAC backfill tables (expected 0 in prod)
  for (const { table, blobCol, hmacCol } of [
    { table: "kyc_customers",    blobCol: "national_id",      hmacCol: "national_id_hmac" },
    { table: "ubo_registrations", blobCol: "owner_national_id", hmacCol: "owner_national_id_hmac" },
  ]) {
    phase2[table] = {};
    const res = await pool.query(
      `SELECT id, ${blobCol}, tenant_id FROM ${table}
       WHERE ${hmacCol} IS NULL AND ${blobCol} IS NOT NULL`
    );
    for (const row of res.rows) {
      phase2[table][row.id as string] = {
        blobValue: row[blobCol] as string,
        tenantId: (row["tenant_id"] as string | null) ?? null,
      };
    }
    const count = Object.keys(phase2[table]).length;
    console.log(`[t004b] [phase-0] ${table} (HMAC backfill): ${count} row(s) need HMAC`);
  }

  return { phase1, phase2 };
}

// ─── Phase 1 — Per-row migration ─────────────────────────────────────────────

async function phase1Migrate(pool: pg.Pool, snap: Phase1Snap): Promise<void> {
  for (const cfg of PHASE1_TABLES) {
    const tableSnap = snap[cfg.table];
    const rowIds = Object.keys(tableSnap);
    let tableMigrated = 0;
    let tableSkipped = 0;

    const { sql: updateSQL, paramBuilder } = buildUpdateSQL(cfg.table, cfg.piiCols);

    for (const rowId of rowIds) {
      const rowSnap = tableSnap[rowId];
      const ctx = cfg.contextFn();

      // Idempotency: skip if primary col is already encrypted
      const primaryVal = rowSnap[cfg.primaryCol];
      if (primaryVal !== null && isEncrypted(primaryVal)) {
        tableSkipped++;
        stats.skipped++;
        continue;
      }

      // Encrypt each PII column (null pass-through)
      const ciphertexts: Record<string, string | null> = {};
      for (const col of cfg.piiCols) {
        const pt = rowSnap[col];
        ciphertexts[col] = pt === null ? null : encryptPii(pt, ctx);
      }

      // ─── Gate A: in-memory pre-write verify ───────────────────────────────
      for (const col of cfg.piiCols) {
        const pt = rowSnap[col];
        const ct = ciphertexts[col];
        if (pt === null && ct === null) continue; // NULL pass-through: correct
        if (pt === null || ct === null) {
          throw new MigrationAbortError(
            `[t004b] Gate A fail — row id=${rowId} (table: ${cfg.table}), col=${col}: null/non-null mismatch. Recovery: ${SNAPSHOT_PATH}`
          );
        }
        const check = decryptPii(ct, ctx);
        if (check !== pt) {
          throw new MigrationAbortError(
            `[t004b] Gate A fail — row id=${rowId} (table: ${cfg.table}), col=${col}: decrypt(encrypt(x)) !== x. Recovery: ${SNAPSHOT_PATH}`
          );
        }
      }

      // ─── BEGIN / conditional UPDATE / COMMIT ──────────────────────────────
      const snapVals = cfg.piiCols.map((col) => rowSnap[col]);
      const cipherVals = cfg.piiCols.map((col) => ciphertexts[col]);
      const params = paramBuilder(rowId, snapVals, cipherVals);

      const client = await pool.connect();
      let committed = false;
      try {
        await client.query("BEGIN");
        const upd = await client.query(updateSQL, params);

        if ((upd.rowCount ?? 0) === 0) {
          await client.query("ROLLBACK");
          throw new MigrationAbortError(
            `[t004b] Stale-overwrite prevented — row id=${rowId} (table: ${cfg.table}) changed since snapshot. Migration aborted. Recovery: ${SNAPSHOT_PATH}`
          );
        }

        await client.query("COMMIT");
        committed = true;
      } catch (err) {
        if (!committed) {
          try { await client.query("ROLLBACK"); } catch (_) { /* ignore */ }
        }
        throw err;
      } finally {
        client.release();
      }

      // ─── Gate B: post-write verify against snapshot ────────────────────────
      // Row is now committed. Read back and verify all columns decrypt correctly.
      const colList = ["id", ...cfg.piiCols].join(", ");
      const readRes = await pool.query(
        `SELECT ${colList} FROM ${cfg.table} WHERE id = $1`,
        [rowId]
      );

      if (readRes.rows.length === 0) {
        throw new MigrationAbortError(
          `[t004b] Gate B fail — row id=${rowId} (table: ${cfg.table}): row not found after commit. Recovery: ${SNAPSHOT_PATH}`
        );
      }

      const dbRow = readRes.rows[0];
      for (const col of cfg.piiCols) {
        const snapVal = rowSnap[col];
        const dbVal = dbRow[col] as string | null;

        if (snapVal === null) {
          if (dbVal !== null) {
            throw new MigrationAbortError(
              `[t004b] Gate B fail — row id=${rowId} (table: ${cfg.table}), col=${col}: expected NULL in DB, got non-NULL. Recovery: ${SNAPSHOT_PATH}`
            );
          }
          continue;
        }
        if (dbVal === null) {
          throw new MigrationAbortError(
            `[t004b] Gate B fail — row id=${rowId} (table: ${cfg.table}), col=${col}: expected ciphertext in DB, got NULL. Recovery: ${SNAPSHOT_PATH}`
          );
        }
        const decrypted = decryptPii(dbVal, ctx);
        if (decrypted !== snapVal) {
          throw new MigrationAbortError(
            `[t004b] Gate B fail — row id=${rowId} (table: ${cfg.table}), col=${col}: decrypted DB value !== snapshot plaintext. Recovery: ${SNAPSHOT_PATH}`
          );
        }
      }

      console.log(`[t004b] confirmed row id=${rowId} (table: ${cfg.table})`);
      tableMigrated++;
      stats.migrated++;
    }

    console.log(
      `[t004b] ${cfg.table}: ${tableMigrated} row(s) migrated+confirmed, ` +
      `${tableSkipped} row(s) skipped (already encrypted)`
    );
  }
}

// ─── Phase 2 — HMAC backfill ─────────────────────────────────────────────────

async function phase2HmacBackfill(pool: pg.Pool, snap: Phase2Snap): Promise<void> {
  for (const { table, blobCol, hmacCol, hmacColName } of [
    {
      table:      "kyc_customers",
      blobCol:    "national_id",
      hmacCol:    "national_id_hmac",
      hmacColName: "national_id",    // column name passed to computeLookupHash
    },
    {
      table:      "ubo_registrations",
      blobCol:    "owner_national_id",
      hmacCol:    "owner_national_id_hmac",
      hmacColName: "owner_national_id",
    },
  ]) {
    const tableSnap = snap[table];
    const rowIds = Object.keys(tableSnap);

    if (rowIds.length === 0) {
      console.log(`[t004b] [phase-2] ${table}: 0 rows HMAC-backfilled (expected 0 in prod)`);
      continue;
    }

    let backfilled = 0;
    console.log(`[t004b] [phase-2] ${table}: ${rowIds.length} row(s) need HMAC backfill`);

    for (const rowId of rowIds) {
      const { blobValue, tenantId } = tableSnap[rowId];

      if (!tenantId) {
        // Pre-FU-015 artifact: tenant_id column was added nullable; legacy rows may have NULL.
        // Cannot compute HMAC without tenant_id (required for per-tenant key derivation).
        // WARN + skip — manual remediation: set tenant_id on this row, then re-run.
        console.warn(
          `[t004b] [phase-2] WARN — row id=${rowId} (table: ${table}): tenant_id IS NULL — ` +
          `HMAC cannot be derived. Row skipped. Manual remediation required: set tenant_id then re-run.`
        );
        stats.nullTenantSkipped++;
        continue;
      }

      const tenantCtx: PiiContext = { type: "tenant", tenantId };

      if (isEncrypted(blobValue)) {
        // Blob already encrypted (e.g. by T004a forward path before T004b ran)
        // Decrypt to get plaintext, compute HMAC, update HMAC only.
        const plaintext = decryptPii(blobValue, tenantCtx);
        const hmac = computeLookupHash(plaintext, tenantId, table, hmacColName);

        const client = await pool.connect();
        try {
          await client.query("BEGIN");
          const upd = await client.query(
            `UPDATE ${table} SET ${hmacCol} = $1 WHERE id = $2 AND ${hmacCol} IS NULL`,
            [hmac, rowId]
          );
          if ((upd.rowCount ?? 0) === 0) {
            await client.query("ROLLBACK");
            // Concurrent update already set the HMAC — skip gracefully
            console.log(`[t004b] [phase-2] row id=${rowId} (table: ${table}): HMAC already set by concurrent update, skipped`);
          } else {
            await client.query("COMMIT");
            console.log(`[t004b] [phase-2] hmac-written row id=${rowId} (table: ${table}, blob was already encrypted)`);
            backfilled++;
            stats.hmacBackfilled++;
          }
        } catch (err) {
          try { await client.query("ROLLBACK"); } catch (_) { /* ignore */ }
          throw err;
        } finally {
          client.release();
        }
      } else {
        // Blob is plaintext — encrypt it AND compute HMAC, update both atomically.
        // blobValue here is the Phase 0 snapshot value (plaintext).
        const plaintext = blobValue;
        const ciphertext = encryptPii(plaintext, tenantCtx);

        // Gate A verify
        const check = decryptPii(ciphertext, tenantCtx);
        if (check !== plaintext) {
          throw new MigrationAbortError(
            `[t004b] [phase-2] Gate A fail — row id=${rowId} (table: ${table}): decrypt(encrypt(x)) !== x. Recovery: ${SNAPSHOT_PATH}`
          );
        }

        const hmac = computeLookupHash(plaintext, tenantId, table, hmacColName);

        // Conditional UPDATE: blobCol must still match snapshot value AND hmacCol still NULL
        const client = await pool.connect();
        let committed = false;
        try {
          await client.query("BEGIN");
          const upd = await client.query(
            `UPDATE ${table} SET ${blobCol} = $1, ${hmacCol} = $2
             WHERE id = $3
               AND ${blobCol} IS NOT DISTINCT FROM $4
               AND ${hmacCol} IS NULL`,
            [ciphertext, hmac, rowId, plaintext]
          );
          if ((upd.rowCount ?? 0) === 0) {
            await client.query("ROLLBACK");
            throw new MigrationAbortError(
              `[t004b] [phase-2] Stale-overwrite prevented — row id=${rowId} (table: ${table}) changed since snapshot. Recovery: ${SNAPSHOT_PATH}`
            );
          }
          await client.query("COMMIT");
          committed = true;
        } catch (err) {
          if (!committed) {
            try { await client.query("ROLLBACK"); } catch (_) { /* ignore */ }
          }
          throw err;
        } finally {
          client.release();
        }

        // Gate B: read back and verify blob decrypts correctly
        const readRes = await pool.query(
          `SELECT ${blobCol} FROM ${table} WHERE id = $1`,
          [rowId]
        );
        if (readRes.rows.length === 0) {
          throw new MigrationAbortError(
            `[t004b] [phase-2] Gate B fail — row id=${rowId} (table: ${table}): row not found after commit. Recovery: ${SNAPSHOT_PATH}`
          );
        }
        const dbBlob = readRes.rows[0][blobCol] as string | null;
        if (!dbBlob || decryptPii(dbBlob, tenantCtx) !== plaintext) {
          throw new MigrationAbortError(
            `[t004b] [phase-2] Gate B fail — row id=${rowId} (table: ${table}): blob decrypt mismatch after commit. Recovery: ${SNAPSHOT_PATH}`
          );
        }

        console.log(`[t004b] [phase-2] hmac-backfilled row id=${rowId} (table: ${table}, blob encrypted+hmac written)`);
        backfilled++;
        stats.hmacBackfilled++;
      }
    }

    console.log(`[t004b] [phase-2] ${table}: ${backfilled} row(s) HMAC-backfilled`);
  }
}

// ─── Main ─────────────────────────────────────────────────────────────────────

async function main(): Promise<void> {
  console.log(`[t004b] ═══════════════════════════════════════════════════`);
  console.log(`[t004b] T004b migration — v1.3-RATIFIED`);
  console.log(`[t004b] DRY_RUN=${DRY_RUN}${DRY_RUN ? " — snapshot written, NO DB writes" : " — LIVE WRITE MODE"}`);
  console.log(`[t004b] ═══════════════════════════════════════════════════`);

  // PROD_DATABASE_URL takes precedence when set (prod execution from dev shell).
  // Falls back to DATABASE_URL (dev execution). Value never logged per Rule 1.
  const connectionString = process.env.PROD_DATABASE_URL ?? process.env.DATABASE_URL;
  const targetEnv = process.env.PROD_DATABASE_URL ? "PROD (via PROD_DATABASE_URL)" : "DEV (via DATABASE_URL)";
  if (!connectionString) {
    throw new Error("[t004b] FATAL: Neither PROD_DATABASE_URL nor DATABASE_URL is set.");
  }
  console.log(`[t004b] Target environment: ${targetEnv}`);

  if (!process.env.PII_AES_MASTER_KEY || process.env.PII_AES_MASTER_KEY.length < 32) {
    throw new Error("[t004b] FATAL: PII_AES_MASTER_KEY not set or < 32 chars.");
  }
  if (!process.env.PII_HMAC_MASTER_KEY || process.env.PII_HMAC_MASTER_KEY.length < 32) {
    throw new Error("[t004b] FATAL: PII_HMAC_MASTER_KEY not set or < 32 chars.");
  }

  const pool = new pg.Pool({ connectionString });

  try {
    // ─── Phase 0: Build and persist encrypted snapshot ─────────────────────
    const snapshot = await phase0BuildSnapshot(pool);

    const p1Total = Object.values(snapshot.phase1)
      .reduce((s, t) => s + Object.keys(t).length, 0);
    const p2Total = Object.values(snapshot.phase2)
      .reduce((s, t) => s + Object.keys(t).length, 0);

    // Encrypt entire snapshot and write to disk (always, even in DRY_RUN)
    const snapshotJson = JSON.stringify(snapshot);
    const encryptedSnapshot = encryptPii(snapshotJson, PLATFORM_CTX);
    fs.writeFileSync(SNAPSHOT_PATH, encryptedSnapshot, "utf8");
    console.log(
      `[t004b] [phase-0] Encrypted snapshot written to ${SNAPSHOT_PATH} ` +
      `— ${p1Total} Phase-1 rows, ${p2Total} Phase-2 rows`
    );

    if (DRY_RUN) {
      for (const [table, rows] of Object.entries(snapshot.phase1)) {
        const ids = Object.keys(rows);
        if (ids.length > 0) {
          console.log(`[t004b] [dry-run] ${table}: ${ids.length} row(s) would be migrated (ids: [${ids.join(", ")}])`);
        }
      }
      for (const [table, rows] of Object.entries(snapshot.phase2)) {
        const ids = Object.keys(rows);
        if (ids.length > 0) {
          console.log(`[t004b] [dry-run] ${table} HMAC backfill: ${ids.length} row(s) (ids: [${ids.join(", ")}])`);
        }
      }
      console.log(`[t004b] [dry-run] No DB writes. Operator checklist before DRY_RUN=false:`);
      console.log(`[t004b] [dry-run]   1. Snapshot file exists at: ${SNAPSHOT_PATH}`);
      console.log(`[t004b] [dry-run]   2. Row count matches expectations (see log above)`);
      console.log(`[t004b] [dry-run]   3. File copied to durable storage (not just this shell)`);
      console.log(`[t004b] [dry-run]   4. PII_AES_MASTER_KEY accessible for recovery-from-snapshot if needed`);
      return;
    }

    // ─── Phase 1: Per-row encryption ───────────────────────────────────────
    console.log(`[t004b] [phase-1] Beginning Phase 1 — ${p1Total} rows`);
    await phase1Migrate(pool, snapshot.phase1);

    // ─── Phase 2: HMAC backfill ────────────────────────────────────────────
    console.log(`[t004b] [phase-2] Beginning Phase 2 — HMAC backfill`);
    await phase2HmacBackfill(pool, snapshot.phase2);

    // ─── Summary ───────────────────────────────────────────────────────────
    console.log(`[t004b] ─── SUMMARY ─────────────────────────────────────`);
    console.log(`[t004b] Phase 1 migrated+confirmed : ${stats.migrated}`);
    console.log(`[t004b] Phase 1 skipped (enc)      : ${stats.skipped}`);
    console.log(`[t004b] Phase 2 HMAC backfilled    : ${stats.hmacBackfilled} (expected: 0 in prod)`);
    console.log(`[t004b] Phase 2 null-tenant skipped: ${stats.nullTenantSkipped} (WARN — manual remediation needed if > 0)`);
    console.log(`[t004b] Snapshot file              : ${SNAPSHOT_PATH}`);
    console.log(`[t004b] DRY_RUN                    : false (live run)`);
    console.log(`[t004b] ─────────────────────────────────────────────────`);

  } finally {
    await pool.end();
  }
}

main().catch((err: Error) => {
  console.error(`[t004b] FATAL: ${err.message}`);
  process.exit(1);
});