import { describe, test, expect, vi } from 'vitest'
import { drainNode, drainHost, DRAIN_REASON } from '../data-plane/drain.js'
import type { AgentTunnel, RouteRegistrar } from '../data-plane/agent-listener.js'
import type { MuxSession } from '../mux/mux-session.js'

function fakeTunnel(hostId: string): AgentTunnel & { drainSpy: ReturnType<typeof vi.fn> } {
  const drainSpy = vi.fn()
  const session = { openStream: vi.fn(), onWire: vi.fn(), drain: drainSpy, close: vi.fn() } satisfies MuxSession
  return { hostId, accountId: 'a', session, closeTunnel: vi.fn(), drainSpy }
}

function registrar(): RouteRegistrar {
  return { register: vi.fn(async () => {}), heartbeat: vi.fn(async () => {}), deregister: vi.fn(async () => {}) }
}

describe('graceful drain + GOAWAY (T11, INV7/INV12)', () => {
  test('operatorDrain: GOAWAY sent, teardown deferred to the grace window', async () => {
    const t = fakeTunnel('h1')
    const reg = registrar()
    const scheduled: (() => void)[] = []
    await drainNode({
      tunnels: () => new Map([['h1', t]]),
      registrar: reg,
      reason: DRAIN_REASON.operatorDrain,
      inFlightGraceMs: 5000,
      setTimeoutFn: (fn) => scheduled.push(fn),
    })
    expect(t.drainSpy).toHaveBeenCalledWith(0, DRAIN_REASON.operatorDrain) // GOAWAY
    expect(t.closeTunnel).not.toHaveBeenCalled() // still within grace
    scheduled.forEach((f) => f()) // grace elapses
    await new Promise((r) => setTimeout(r, 0))
    expect(t.closeTunnel).toHaveBeenCalled()
    expect(reg.deregister).toHaveBeenCalledWith('h1')
  })

  test('revoked: IMMEDIATE teardown regardless of a large grace (Finding-2/INV12)', async () => {
    const t = fakeTunnel('h1')
    const reg = registrar()
    const setTimeoutFn = vi.fn()
    await drainHost('h1', {
      tunnels: () => new Map([['h1', t]]),
      registrar: reg,
      reason: DRAIN_REASON.revoked,
      inFlightGraceMs: 30_000, // must be ignored
      setTimeoutFn,
    })
    expect(t.drainSpy).toHaveBeenCalledWith(0, DRAIN_REASON.revoked)
    expect(t.closeTunnel).toHaveBeenCalled() // immediate
    expect(reg.deregister).toHaveBeenCalledWith('h1')
    expect(setTimeoutFn).not.toHaveBeenCalled() // grace never scheduled
  })

  test('drainHost targets exactly one host; siblings keep running', async () => {
    const h1 = fakeTunnel('h1')
    const h2 = fakeTunnel('h2')
    const reg = registrar()
    await drainHost('h1', {
      tunnels: () => new Map([['h1', h1], ['h2', h2]]),
      registrar: reg,
      reason: DRAIN_REASON.revoked,
      inFlightGraceMs: 0,
    })
    expect(h1.closeTunnel).toHaveBeenCalled()
    expect(h2.closeTunnel).not.toHaveBeenCalled()
    expect(reg.deregister).toHaveBeenCalledWith('h1')
    expect(reg.deregister).not.toHaveBeenCalledWith('h2')
  })

  test('drainHost for an absent host is a no-op (deny-by-default, never a broader kill)', async () => {
    const reg = registrar()
    await drainHost('ghost', { tunnels: () => new Map(), registrar: reg, reason: DRAIN_REASON.revoked, inFlightGraceMs: 0 })
    expect(reg.deregister).not.toHaveBeenCalled()
  })
})