import assert from 'node:assert/strict'
import { spawnSync } from 'node:child_process'
import { readFileSync } from 'node:fs'
import { describe, it } from 'node:test'
import { parseProductionCapacityCellArguments } from './prepare-relay-production-capacity-canary.mjs'
import {
  SAME_CAP_CELLS,
  SAME_CAP_MIGRATION_ONLY_CELLS,
  entryAdmission,
  selectorWaveDelta
} from './relay-production-same-cap-wave.mjs'
import { readRelayWorkflow } from './relay-repository.mjs'
import { validateCapacityPlan } from './validate-relay-capacity-plan.mjs'

const workflow = readRelayWorkflow('deploy-relay-production-same-cap-job.yml')
const capacityWorkflow = readRelayWorkflow('deploy-relay-production-capacity-job.yml')
const production = readFileSync(
  new URL('../../infra/terraform/environments/production.tfvars', import.meta.url),
  'utf8'
)
const REHOME_SOURCE_CELLS = rehomeSourceCells()
const DIRECTOR_IDENTITY = 'relay-director@onorca-cloud.iam.gserviceaccount.com'
const AUDIENCE = 'https://relay.onorca.dev/v1/admin/host-drain'
const ROLLBACK_IMAGE = `us-central1-docker.pkg.dev/p/orca-cloud/relay@sha256:${'d'.repeat(64)}`
const TARGET_IMAGE = `us-central1-docker.pkg.dev/p/orca-cloud/relay@sha256:${'e'.repeat(64)}`

// The startup template emits rehome trust only for cells in this list, so it is what decides
// whether a cell's plan may carry those lines at all.
function rehomeSourceCells() {
  const start = production.indexOf('relay_region_rehome_source_cell_ids = [')
  assert.notEqual(start, -1, 'production.tfvars has no rehome source cell list')
  const end = production.indexOf(']', start)
  assert.notEqual(end, -1, 'the rehome source cell list is unterminated')
  return new Set(
    [...production.slice(start, end).matchAll(/"([^"]+)"/g)].map(([, cell]) => cell)
  )
}

// The job cross-checks its pinned pool against the committed map; model the same read.
function tfvarsDatabasePoolMax(cellId) {
  return tfvarsCellBlock(cellId).match(/database_pool_max\s*=\s*(\d+)/)?.[1] ?? '10'
}

function tfvarsHardCap(cellId) {
  const cap = /connection_hard_cap\s*=\s*(\d+)/.exec(tfvarsCellBlock(cellId))?.[1]
  assert.notEqual(cap, undefined, `${cellId} has no connection_hard_cap`)
  return cap
}

function tfvarsCellBlock(cellId) {
  const start = production.indexOf(`"${cellId}" = {`)
  assert.notEqual(start, -1, `${cellId} is missing from production.tfvars`)
  return production.slice(start, production.indexOf('\n  }', start))
}

function startupScript({ cap, image, trusted, pool }) {
  return [
    `  printf 'ORCA_RELAY_CELL_CONNECTION_HARD_CAP=%s\\n' '${cap}'`,
    `  printf 'ORCA_RELAY_CELL_CONNECTION_UNOBSERVED_BOUND=%s\\n' '60'`,
    ...(pool === undefined
      ? []
      : [`  printf 'ORCA_RELAY_DATABASE_POOL_MAX=%s\\n' '${pool}'`]),
    ...(trusted ? [
      `  printf 'ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT=%s\\n' '${DIRECTOR_IDENTITY}'`,
      `  printf 'ORCA_RELAY_REHOME_AUDIENCE=%s\\n' '${AUDIENCE}'`
    ] : []),
    `printf 'ORCA_RELAY_IMAGE_DIGEST=%s\\n' '${image.split('@')[1]}'`,
    `docker pull '${image}'`,
    'docker run --detach \\',
    '  --name orca-relay \\',
    `  '${image}'`
  ].join('\n')
}

// The exact shape the apply step's plan has: template replaced, MIG rebound to it.
function rollPlan({ cellId, cap, protocol, pool }) {
  return {
    configuration: {
      root_module: {
        resources: [{
          address: 'google_compute_instance_group_manager.relay_gce_cell',
          expressions: {
            version: [{
              instance_template: {
                references: [
                  'google_compute_instance_template.relay_gce_cell',
                  'each.key'
                ]
              },
              name: { constant_value: 'primary' }
            }]
          }
        }]
      }
    },
    resource_changes: [
      {
        address: `google_compute_instance_template.relay_gce_cell[${JSON.stringify(cellId)}]`,
        change: {
          actions: ['create', 'delete'],
          before: {
            metadata_startup_script: startupScript({
              cap,
              image: ROLLBACK_IMAGE,
              trusted: protocol >= 1,
              // The live template predates the reviewed pool raise, as every asia cell's does.
              pool: pool === undefined ? undefined : '10'
            })
          },
          after: {
            metadata_startup_script: startupScript({
              cap,
              image: TARGET_IMAGE,
              trusted: protocol >= 1,
              pool
            }),
            self_link: null
          },
          after_unknown: { self_link: true }
        }
      },
      {
        address: `google_compute_instance_group_manager.relay_gce_cell[${JSON.stringify(cellId)}]`,
        change: {
          actions: ['update'],
          before: { target_size: 1, version: [{ instance_template: 'old' }] },
          after: { target_size: 1, version: [{ instance_template: null }] },
          after_unknown: { version: [{ instance_template: true }] }
        }
      }
    ]
  }
}

function hostname(cellId) {
  return cellId.slice('production-gce-'.length)
}

// The job resolves cap, region, and pool from the cell id before any admin call; run that block
// alone. An empty pool is the root default, which the startup template emits no line for.
function resolveCellShape(cellId) {
  const start = workflow.indexOf('          TARGET_HOSTNAME="${TARGET_CELL_ID#production-gce-}"')
  assert.notEqual(start, -1, 'the same-cap cell shape block is missing')
  const end = workflow.indexOf('\n          esac\n', start)
  assert.notEqual(end, -1, 'the same-cap cell shape block has no esac')
  const script = workflow.slice(start, end + '\n          esac'.length).replace(/^ {10}/gm, '')
  return spawnSync('bash', [
    '-euo',
    'pipefail',
    '-c',
    `${script}\necho "\${EXPECTED_REGION} \${EXPECTED_HARD_CAP} pool=\${EXPECTED_DATABASE_POOL_MAX}"`
  ], { env: { ...process.env, TARGET_CELL_ID: cellId }, encoding: 'utf8' })
}

function cellShape(cellId) {
  const resolved = resolveCellShape(cellId)
  assert.equal(resolved.status, 0, `${cellId}: ${resolved.stderr}`)
  const [, cap, pool] = resolved.stdout.trim().split(' ')
  return { cap: Number(cap), pool: pool.slice('pool='.length) || undefined }
}

// The class block runs before checkout-independent work and decides the whole wave shape.
function resolveCellClass(cellId) {
  return spawnSync('bash', [
    '-euo',
    'pipefail',
    '-c',
    `${jobBlock(
      '          CELL_CLASS="$(node dev/scripts/relay-production-same-cap-wave.mjs cell-class \\',
      '          SELECTOR_WAVE_DELTA="$(jq -er \'.selectorWaveDelta\' <<< "${CELL_CLASS}")"'
    )}\necho "\${ENTRY_ADMISSION} \${SELECTOR_WAVE_DELTA}"`
  ], { cwd: new URL('../..', import.meta.url), env: { ...process.env, TARGET_CELL_ID: cellId }, encoding: 'utf8' })
}

function generationBlock() {
  return `${jobBlock(
    '          if test "${DEPLOY_MODE}" = verify; then',
    '          fi'
  )}\necho "\${EFFECTIVE_SELECTOR_GENERATION}"`
}

// The job derives both memberships in one block; run that block alone for each class.
function membership(env) {
  const script = `${jobBlock(
    '          RESTORED_MIGRATION_CELLS="$(jq -rn \\',
    '          fi'
  )}\njq -cn --arg a "\${ISOLATED_MIGRATION_CELLS}" --arg b "\${ISOLATED_GENERAL_CELLS}" \\
  --arg c "\${RESTORED_MIGRATION_CELLS}" --arg d "\${RESTORED_GENERAL_CELLS}" \\
  '{isolatedMigration:$a,isolatedGeneral:$b,restoredMigration:$c,restoredGeneral:$d}'`
  const resolved = spawnSync('bash', ['-euo', 'pipefail', '-c', script], {
    env: { ...process.env, ...env },
    encoding: 'utf8'
  })
  assert.equal(resolved.status, 0, resolved.stderr)
  return JSON.parse(resolved.stdout)
}

function jobBlock(firstLine, lastLine) {
  const start = workflow.indexOf(`${firstLine}\n`)
  assert.notEqual(start, -1, `the job has no ${firstLine.trim()}`)
  const end = workflow.indexOf(`\n${lastLine}\n`, start)
  assert.notEqual(end, -1, `that block has no ${lastLine.trim()}`)
  return workflow.slice(start, end + lastLine.length + 1).replace(/^ {10}/gm, '')
}

describe('same-cap roll scripts accept every same-cap cell', () => {
  it('parses every wave cell through the same-cap canary allowlist', () => {
    for (const cellId of SAME_CAP_CELLS) {
      for (const mode of ['isolate', 'drain', 'activate']) {
        assert.deepEqual(parseProductionCapacityCellArguments([
          '--director-origin', 'https://relay.onorca.dev',
          '--cell-origin', `https://${hostname(cellId)}.relay.onorca.dev`,
          '--cell-id', cellId,
          '--approved-cells', 'same-cap',
          '--mode', mode
        ]), {
          directorOrigin: 'https://relay.onorca.dev',
          cellOrigin: `https://${hostname(cellId)}.relay.onorca.dev`,
          cellId,
          mode,
          paceWindowMs: 0
        })
      }
    }
  })

  it('resolves a cap, region, and pool for every wave cell and refuses anything else', () => {
    for (const cellId of SAME_CAP_CELLS) {
      const resolved = resolveCellShape(cellId)
      assert.equal(resolved.status, 0, `${cellId}: ${resolved.stderr}`)
      assert.match(
        resolved.stdout.trim(),
        /^(us-central1 1000 pool=|us-central1 600 pool=|asia-east2 3000 pool=16)$/,
        cellId
      )
      assert.equal(tfvarsDatabasePoolMax(cellId), cellShape(cellId).pool ?? '10', cellId)
      assert.equal(String(cellShape(cellId).cap), tfvarsHardCap(cellId), cellId)
    }
    assert.equal(resolveCellShape('production-gce-c12').status, 1)
    assert.equal(resolveCellShape('production-gce-c30').status, 1)
  })

  it('passes the same-cap allowlist on every canary invocation the job runs', () => {
    const invocations = workflow.split('prepare-relay-production-capacity-canary.mjs').slice(1)
    assert.equal(invocations.length, 4)
    for (const invocation of invocations) {
      const lines = invocation.split('\n')
      const end = lines.findIndex((line) => !line.endsWith('\\'))
      const call = lines.slice(0, end + 1).join(' ')
      assert.match(call, /--approved-cells same-cap/)
      // The restore call picks its mode from the cell's entry admission class.
      assert.match(call, /--mode (isolate|drain|activate|"\$\{RESTORE_MODE\}")/)
    }
  })

  it('paces the drain it sends to the selected cell', () => {
    const drain = workflow.split('--mode drain')[1] ?? ''
    assert.match(drain.split('\n').slice(0, 2).join(' '), /--pace-window-ms "\$\{DRAIN_PACE_WINDOW_MS\}"/)
    assert.match(workflow, /DRAIN_PACE_WINDOW_MS: '120000'/)
    // The transition wait has to outlast the pacing window on top of the leases it waits on.
    assert.match(workflow, /--activity restart-safe[\s\S]*?--timeout-ms 1020000/)
  })

  it('passes this cell\'s rehome protocol and pool on every plan validation the job runs', () => {
    const invocations = workflow.split('validate-relay-capacity-plan.mjs').slice(1)
    assert.equal(invocations.length, 2)
    for (const invocation of invocations) {
      const lines = invocation.split('\n')
      const end = lines.findIndex((line) => !line.trimEnd().endsWith('\\'))
      const call = lines.slice(0, end + 1).join(' ')
      assert.match(call, /--mode same-cap-cell/)
      assert.match(call, /--regional-rehome-protocol "\$\{DESIRED_REHOME_PROTOCOL\}"/)
      assert.match(call, /"\$\{POOL_ARGUMENTS\[@\]\}"/)
    }
    // Each of those steps must build the flag from the resolved pool, and only when there is one.
    const builders = workflow.split(
      'if test -n "${EXPECTED_DATABASE_POOL_MAX}"; then\n' +
        '            POOL_ARGUMENTS=(--database-pool-max "${EXPECTED_DATABASE_POOL_MAX}")'
    )
    assert.equal(builders.length, 3)
    assert.equal(workflow.split('POOL_ARGUMENTS=()').length, 3)
  })

  it('validates a correct plan for every wave cell at that cell\'s rehome protocol', () => {
    const trusted = SAME_CAP_CELLS.filter((cell) => REHOME_SOURCE_CELLS.has(cell))
    // Only a declared rehome source may roll at a trusted protocol at all; the job refuses
    // the rest before it plans, and the next test covers them at protocol 0.
    assert.deepEqual(
      SAME_CAP_CELLS.filter((cell) => !REHOME_SOURCE_CELLS.has(cell)),
      SAME_CAP_MIGRATION_ONLY_CELLS
    )
    for (const [cellId, protocol] of trusted.flatMap((cell) => [[cell, 1], [cell, 3]])) {
      const { cap, pool } = cellShape(cellId)
      const config = {
        mode: 'same-cap-cell',
        cellId,
        hardCap: cap,
        unobservedBound: 60,
        image: TARGET_IMAGE,
        rollbackImage: ROLLBACK_IMAGE,
        rehomeDirectorServiceAccount: DIRECTOR_IDENTITY,
        rehomeAudience: AUDIENCE,
        regionalRehomeProtocol: String(protocol),
        databasePoolMax: pool
      }
      const plan = rollPlan({ cellId, cap, protocol, pool })
      assert.deepEqual(
        validateCapacityPlan(plan, config),
        { mode: 'same-cap-cell', changes: 2 },
        cellId
      )
      // The other protocol must reject the same plan, or the flag decides nothing.
      assert.throws(
        () => validateCapacityPlan(plan, {
          ...config,
          regionalRehomeProtocol: '0'
        }),
        /reviewed image and capacity/,
        cellId
      )
      // Dropping the pin must reject a pinned cell, and adding one must reject a default cell.
      assert.throws(
        () => validateCapacityPlan(plan, {
          ...config,
          databasePoolMax: pool === undefined ? '16' : undefined
        }),
        /reviewed image and capacity/,
        cellId
      )
    }
  })

  it('validates a protocol-0 plan for a cell outside the rehome source list', () => {
    const cellId = 'production-gce-c17'
    assert.equal(REHOME_SOURCE_CELLS.has(cellId), false)
    const config = {
      mode: 'same-cap-cell',
      cellId,
      hardCap: 600,
      unobservedBound: 60,
      image: TARGET_IMAGE,
      rollbackImage: ROLLBACK_IMAGE,
      rehomeDirectorServiceAccount: DIRECTOR_IDENTITY,
      rehomeAudience: AUDIENCE,
      regionalRehomeProtocol: '0'
    }
    const plan = rollPlan({ cellId, cap: 600, protocol: 0 })
    assert.deepEqual(validateCapacityPlan(plan, config), { mode: 'same-cap-cell', changes: 2 })
    // Protocol 1 must reject a plan with no rehome lines, or the absent-line rule decides nothing.
    assert.throws(
      () => validateCapacityPlan(plan, { ...config, regionalRehomeProtocol: '1' }),
      /reviewed image and capacity/
    )
  })

  it('resolves the class and selector delta the wave validator declares', () => {
    for (const cellId of SAME_CAP_CELLS) {
      const resolved = resolveCellClass(cellId)
      assert.equal(resolved.status, 0, `${cellId}: ${resolved.stderr}`)
      assert.equal(
        resolved.stdout.trim(),
        `${entryAdmission(cellId)} ${selectorWaveDelta(cellId)}`,
        cellId
      )
    }
    assert.equal(resolveCellClass('production-gce-c12').status, 1)
  })

  it('offsets a later wave by this cell class\'s own selector delta', () => {
    for (const [waveIndex, delta] of [['0', 2], ['3', 2], ['0', 0], ['3', 0]]) {
      const resolved = spawnSync('bash', ['-euo', 'pipefail', '-c', generationBlock()], {
        env: {
          ...process.env,
          DEPLOY_MODE: 'apply',
          EXPECTED_SELECTOR_GENERATION: '40',
          WAVE_INDEX: waveIndex,
          SELECTOR_WAVE_DELTA: String(delta)
        },
        encoding: 'utf8'
      })
      assert.equal(resolved.status, 0, resolved.stderr)
      assert.equal(resolved.stdout.trim(), String(40 + delta * Number(waveIndex)))
    }
  })

  it('hands a migration-only cell back the exact membership it entered with', () => {
    const entry = {
      EXPECTED_MIGRATION_ONLY_CELLS: 'production-gce-c17,production-gce-c18',
      EXPECTED_GENERAL_CELLS: 'production-gce-c7,production-gce-c8'
    }
    const isolated = membership({
      ...entry,
      TARGET_CELL_ID: 'production-gce-c17',
      ENTRY_ADMISSION: 'migration-only'
    })
    assert.deepEqual(isolated, {
      isolatedMigration: 'production-gce-c17,production-gce-c18',
      isolatedGeneral: 'production-gce-c7,production-gce-c8',
      restoredMigration: 'production-gce-c17,production-gce-c18',
      restoredGeneral: 'production-gce-c7,production-gce-c8'
    })
    // A general cell still leaves migration-only and returns to general.
    assert.deepEqual(
      membership({
        ...entry,
        TARGET_CELL_ID: 'production-gce-c7',
        ENTRY_ADMISSION: 'general'
      }),
      {
        isolatedMigration: 'production-gce-c17,production-gce-c18,production-gce-c7',
        isolatedGeneral: 'production-gce-c8',
        restoredMigration: 'production-gce-c17,production-gce-c18',
        restoredGeneral: 'production-gce-c7,production-gce-c8'
      }
    )
  })

  it('never activates a migration-only cell and proves its isolate changed nothing', () => {
    const restore = workflow
      .split('name: Restore only the verified selected cell to its entry admission')[1]
      .split('\n      - id:')[0]
    assert.match(restore, /if test "\$\{ENTRY_ADMISSION\}" = migration-only; then\n\s+RESTORE_MODE=isolate/)
    assert.match(restore, /--admission "\$\{ENTRY_ADMISSION\}"/)
    // The pre-mutation check must demand the class the cell is declared to serve in.
    assert.match(workflow, /PRECHECK_ADMISSION="\$\{ENTRY_ADMISSION\}"/)
    const isolate = workflow
      .split('name: Reversibly isolate and drain only the selected cell')[1]
      .split('\n      - id:')[0]
    assert.match(isolate, /migration-only; then\n\s+jq -e '\.changed == false'/)
  })

  it('requires rehome source membership exactly when a roll carries trust lines', () => {
    const step = workflow
      .split('name: Resolve immutable same-cap cell configuration')[1]
      .split('\n      - name:')[0]
    const guard = step.indexOf('jq -e --arg cell "${TARGET_CELL_ID}" \'index($cell) != null\'')
    assert.notEqual(guard, -1)
    // The guard reads both protocols, so it has to sit after they are resolved.
    assert.ok(step.indexOf('DESIRED_REHOME_PROTOCOL="${TARGET_REHOME_PROTOCOL}"') < guard)
    assert.match(
      step.slice(0, guard),
      /test "\$\{DESIRED_REHOME_PROTOCOL\}" != 0 \|\| test "\$\{CURRENT_REHOME_PROTOCOL\}" != 0\n\s+\}; then\s+$/
    )
  })

  it('leaves the US-only capacity job on the default allowlist', () => {
    assert.doesNotMatch(capacityWorkflow, /--approved-cells/)
  })
})

// Both trusted versions must prove the same authenticated drain boundary.
it('proves rehome trust for protocol 3 on forward and rollback rolls', () => {
  const step = workflow.split('name: Prove exact per-host trust and idempotent no-neighbor behavior')[1].split('\n      - name:')[0]
  assert.match(step, /inputs\.rollback-rehome-protocol != '0'/)
  assert.match(step, /inputs\.target-rehome-protocol != '0'/)
  assert.match(step, /probe-relay-rehome-trust\.mjs/)
})