import { test } from '@japa/runner'
import db from '@adonisjs/lucid/services/db'
import testUtils from '@adonisjs/core/services/test_utils'
import { mkdtemp, rm } from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import coordinatorConfig from '#config/coordinator'
import GroupLeases from '#services/group_leases'
import Leases, { LeaseLost } from '#services/leases'
import { Readable } from 'node:stream'

async function node(name: string) {
  return Number((await db.table('nodes').insert({ name }).returning('id'))[0].id)
}

async function group(name: string, priority: number) {
  return Number((await db.table('groups').insert({ name, priority }).returning('id'))[0].id)
}

async function range(groupId: number, start: number, priority: number, extra: Record<string, unknown> = {}) {
  const [row] = await db
    .table('ranges')
    .insert({ group_id: groupId, start_id: start, end_id: start + 99, next_id: start, priority, ...extra })
    .returning('id')
  return Number(row.id)
}

const rangeRow = async (id: number) => (await db.from('ranges').where('id', id).firstOrFail()) as any

test.group('Group leases', (g) => {
  let dir = ''
  let snapshotDir = ''
  g.each.setup(async () => {
    const rollback = await testUtils.db().withGlobalTransaction()
    dir = await mkdtemp(join(tmpdir(), 'snapshots-'))
    snapshotDir = coordinatorConfig.snapshotDir
    coordinatorConfig.snapshotDir = dir
    return rollback
  })
  g.each.teardown(async () => {
    coordinatorConfig.snapshotDir = snapshotDir
    await rm(dir, { recursive: true, force: true })
  })

  test('a lease takes the group with the most urgent planned range of its lane; its ranges come in article order; caught up it gives nothing', async ({
    assert,
  }) => {
    const a = await group('alt.binaries.gl.a', 1)
    const b = await group('alt.binaries.gl.b', 2)
    // Group a: history at priority 0 and 1; group b: a tail range (priority -3) and history.
    const a1 = await range(a, 1, 0)
    const a2 = await range(a, 101, 0)
    const a3 = await range(a, 201, 1)
    const b2 = await range(b, 101, 0)
    const b1 = await range(b, 1, -3)
    const n1 = await node('gl-1')
    const n2 = await node('gl-2')

    // The history lane: both groups have priority 0 work; a has the most planned articles (300 to 100).
    const first = await GroupLeases.lease(n1)
    assert.equal(first?.group.id, a, 'the most planned articles first')
    assert.equal(first?.lane, 'history')
    assert.isNull(first?.snapshot)
    // Held: the next node gets the other group; a third gets nothing.
    const second = await GroupLeases.lease(n2)
    assert.equal(second?.group.id, b)
    assert.isNull(await GroupLeases.lease(await node('gl-3')))

    // Ranges in article order, whatever their priority; the tail range isn't the history lane's.
    assert.deepEqual((await GroupLeases.next(b, n2))?.range, { id: b2, start: 101, end: 200, next: 101 })
    assert.include(await rangeRow(b2), { status: 'leased', node_id: n2 })
    // Asked again before finishing: the same range.
    assert.equal((await GroupLeases.next(b, n2))?.range.id, b2)
    // A chunk of b2 that doesn't inflate, reported after the batch before it: the range is asked for again and
    // comes back at its resume point (the chunk fetched plain next time).
    await db.from('ranges').where('id', b2).update({ next_id: 151 })
    await Leases.corrupt(b2, n2, 151)
    assert.include(await rangeRow(b2), { status: 'planned', node_id: null })
    assert.deepEqual((await GroupLeases.next(b, n2))?.range, { id: b2, start: 101, end: 200, next: 151 })
    assert.deepEqual((await rangeRow(b2)).plain_chunks.map(Number), [151])
    await db.from('ranges').where('id', b2).update({ status: 'done', next_id: 201 })
    assert.isNull(await GroupLeases.next(b, n2), 'caught up: the tail range is the other lane\'s')
    assert.include(await rangeRow(b1), { status: 'planned' })

    // The tail lane: group b's tail, held by a third node while n2 holds b's history.
    const n3 = await node('gl-3b')
    const tail = await GroupLeases.lease(n3, 'tail')
    assert.equal(tail?.group.id, b)
    assert.equal(tail?.lane, 'tail')
    assert.isNull(await GroupLeases.lease(await node('gl-3c'), 'tail'), 'no other group has tail work')
    assert.deepEqual((await GroupLeases.next(b, n3, 'tail'))?.range, { id: b1, start: 1, end: 100, next: 1 })
    // The history lane is n2's.
    await assert.rejects(() => GroupLeases.next(b, n3), LeaseLost)
    await db.from('ranges').where('id', b1).update({ status: 'done', next_id: 101 })
    assert.isNull(await GroupLeases.next(b, n3, 'tail'), 'the tail lane is caught up')
    // Each lane its own snapshot.
    await GroupLeases.putSnapshot(b, n3, 101, Readable.from([Buffer.from('tail at 101')]), 'tail')
    await GroupLeases.putSnapshot(b, n2, 201, Readable.from([Buffer.from('history at 201')]))
    assert.equal((await GroupLeases.getSnapshot(b, n3, 'tail'))?.next, 101)
    assert.equal((await GroupLeases.getSnapshot(b, n2))?.next, 201)
    await GroupLeases.release(b, n3, 'tail')
    assert.include(await GroupLeases.held(b, n2), { node_id: n2 })
    await assert.rejects(() => GroupLeases.held(b, n3, 'tail'), LeaseLost)

    assert.deepEqual([a1, a2, a3].length, 3)
    await assert.rejects(() => GroupLeases.next(a, n2), LeaseLost)
  })

  test('the order mixes the group with the most planned articles and the longest unheld one', async ({ assert }) => {
    // Four groups with priority 0 work: c was held before (5000 articles left), a, b and d never (100, 1000, 50).
    const a = await group('alt.binaries.gl.o.a', 1)
    const b = await group('alt.binaries.gl.o.b', 2)
    const c = await group('alt.binaries.gl.o.c', 3)
    const d = await group('alt.binaries.gl.o.d', 4)
    await range(a, 1, 0)
    for (let i = 0; i < 10; i++) await range(b, 1 + 100 * i, 0)
    for (let i = 0; i < 50; i++) await range(c, 1 + 100 * i, 0)
    await db.table('ranges').insert({ group_id: d, start_id: 1, end_id: 50, next_id: 1, priority: 0 })
    await db.table('group_leases').insert({ group_id: c, lane: 'history', node_id: null, heartbeat_at: new Date(Date.now() - 3600_000) })
    // By articles c, b, a, d; by age a, b, d (never held, by group priority), c. The leases alternate which
    // ranking goes first: c, a, b, d (or a, c, b, d when the sequence stands on the other parity).
    await db.rawQuery(`SELECT setval('group_lease_turns', 2)`)
    const order = []
    for (const name of ['o1', 'o2', 'o3', 'o4']) order.push((await GroupLeases.lease(await node(name)))?.group.id)
    assert.deepEqual(order, [c, a, b, d])
  })

  test('a stale holder is taken over and the ranges after the snapshot are read again; a release keeps the resume point', async ({
    assert,
  }) => {
    const a = await group('alt.binaries.gl.c', 1)
    const r1 = await range(a, 1, 0)
    const r2 = await range(a, 101, 0)
    const r3 = await range(a, 201, 0)
    const r4 = await range(a, 301, 0)
    const n1 = await node('gl-4')
    const n2 = await node('gl-5')
    assert.equal((await GroupLeases.lease(n1))?.group.id, a)
    // The worker did r1 and r2, snapshotted at 201 (r3's start), then got halfway through r3 and went quiet.
    await db.from('ranges').whereIn('id', [r1, r2]).update({ status: 'done', next_id: db.raw('end_id + 1') })
    await GroupLeases.putSnapshot(a, n1, 201, Readable.from([Buffer.from('window at 201')]))
    await GroupLeases.next(a, n1)
    await db.from('ranges').where('id', r3).update({ next_id: 251 })
    await db.from('group_leases').where('group_id', a).update({ heartbeat_at: new Date(Date.now() - 24 * 3600_000) })

    const lease = await GroupLeases.lease(n2)
    assert.equal(lease?.group.id, a)
    assert.deepEqual(lease?.snapshot, { next: 201 })
    // r1 and r2 stay done (before the snapshot); r3 is read again from the snapshot's article, r4 untouched.
    assert.include(await rangeRow(r1), { status: 'done' })
    assert.include(await rangeRow(r2), { status: 'done' })
    assert.include(await rangeRow(r3), { status: 'planned', node_id: null })
    assert.equal(Number((await rangeRow(r3)).next_id), 201)
    assert.include(await rangeRow(r4), { status: 'planned' })
    const snapshot = await GroupLeases.getSnapshot(a, n2)
    assert.equal(snapshot?.next, 201)
    assert.equal(snapshot?.size, 'window at 201'.length)
    await assert.rejects(() => GroupLeases.next(a, n1), LeaseLost)

    // The new holder works r3, snapshots at 301 and gives the group back: r3 done, r4 waits at its start.
    assert.equal((await GroupLeases.next(a, n2))?.range.id, r3)
    await db.from('ranges').where('id', r3).update({ status: 'done', next_id: 301 })
    await GroupLeases.putSnapshot(a, n2, 301, Readable.from([Buffer.from('window at 301')]))
    assert.equal((await GroupLeases.next(a, n2))?.range.id, r4)
    await db.from('ranges').where('id', r4).update({ next_id: 301 })
    await GroupLeases.release(a, n2)
    assert.include(await rangeRow(r4), { status: 'planned', node_id: null })
    assert.isNull((await db.from('group_leases').where('group_id', a).firstOrFail()).node_id)
    // Leased again at once (free, not stale): the snapshot at 301 covers everything done.
    const again = await GroupLeases.lease(n1)
    assert.deepEqual(again?.snapshot, { next: 301 })
    assert.include(await rangeRow(r3), { status: 'done' })
    assert.equal((await GroupLeases.next(a, n1))?.range.id, r4)
  })

  test('capture only: the history ranges sharing a window priority are neither leased nor handed out', async ({
    assert,
  }) => {
    // Group e: history at priority 0 only (the planner's oldest bucket); group f: a 2016 window range (capture)
    // between two history ranges at 0.
    const e = await group('alt.binaries.gl.e', 1)
    await range(e, 1, 0)
    const f = await group('alt.binaries.gl.f', 2)
    await range(f, 1, 0)
    const f2 = await range(f, 101, 0, { capture: true })
    await range(f, 201, 0)
    const n1 = await node('gl-8')
    const was = coordinatorConfig.leaseCaptureOnly
    coordinatorConfig.leaseCaptureOnly = true
    try {
      assert.equal((await GroupLeases.lease(n1))?.group.id, f, 'the group with a capture range, not the older one')
      assert.equal((await GroupLeases.next(f, n1))?.range.id, f2)
      await db.from('ranges').where('id', f2).update({ status: 'done', next_id: 201 })
      assert.isNull(await GroupLeases.next(f, n1), 'the history range after the window is not handed out')
      assert.isNull(await GroupLeases.lease(await node('gl-9')))
    } finally {
      coordinatorConfig.leaseCaptureOnly = was
    }
    assert.equal((await GroupLeases.lease(await node('gl-10')))?.group.id, e, 'any range again without the flag')
  })

  test('a snapshot older than a done range makes that range be read again from its start', async ({ assert }) => {
    const a = await group('alt.binaries.gl.d', 1)
    const r1 = await range(a, 1, 0)
    const r2 = await range(a, 101, 0)
    const n1 = await node('gl-6')
    await GroupLeases.lease(n1)
    // A range after the snapshot's article done long before the snapshot was taken (the old range leases did it):
    // not the window's, it stays done.
    const r3 = await range(a, 201, 0)
    await db.from('ranges').where('id', r3).update({ status: 'done', next_id: 301, done_at: new Date(Date.now() - 48 * 3600_000) })
    // The worker snapshotted at the start, finished r1, and crashed before its next snapshot: r1 is read again.
    await GroupLeases.putSnapshot(a, n1, 1, Readable.from([Buffer.from('window at 1')]))
    await db.from('ranges').where('id', r1).update({ status: 'done', next_id: 101, done_at: new Date() })
    await db.from('group_leases').where('group_id', a).update({ heartbeat_at: new Date(Date.now() - 24 * 3600_000) })
    await GroupLeases.lease(await node('gl-7'))
    assert.equal((await rangeRow(r1)).status, 'planned')
    assert.equal(Number((await rangeRow(r1)).next_id), 1)
    assert.equal((await rangeRow(r2)).status, 'planned')
    assert.equal(Number((await rangeRow(r2)).next_id), 101)
    assert.equal((await rangeRow(r3)).status, 'done')
  })
})