import { BaseCommand, flags } from '@adonisjs/core/ace'
import type { CommandOptions } from '@adonisjs/core/types/ace'
import { setTimeout as sleep } from 'node:timers/promises'
import { gzipSync, unzipSync } from 'node:zlib'
import { RangeIndexer, decodeFrames, toWire, type EdgeBatch, type Lease } from '#lib/edge/index'

const IDLE_MS = 30_000
const RETRIES = 8

class LeaseLost extends Error {}

/**
 * An edge worker (REBUILD.md, part 2): leases a range from the coordinator,
 * streams its chunks, indexes them, posts a batch every --batch-chunks
 * chunks (each moves the range's resume point), and finishes the range.
 * Needs COORDINATOR_URL and NODE_TOKEN only. On SIGTERM it gives its range
 * back after the current chunk; the next worker resumes from the last batch.
 */
export default class EdgeNode extends BaseCommand {
  static commandName = 'edge:node'
  static description = 'Indexes ranges for the coordinator'
  static options: CommandOptions = { startApp: true, staysAlive: true }

  @flags.number({ description: 'Chunks per posted batch (500k articles: ~4% of releases get a second chunk; LOG.md 2026-10-09)', default: 10 })
  declare batchChunks: number

  @flags.number({ description: 'Stop after this many ranges (default: keep going)' })
  declare ranges?: number

  private stopping = false

  async run() {
    const url = process.env.COORDINATOR_URL
    const token = process.env.NODE_TOKEN
    if (!url || !token) throw new Error('COORDINATOR_URL and NODE_TOKEN are needed')
    let stopped: () => void
    const done = new Promise<void>((resolve) => (stopped = resolve))
    this.app.terminating(async () => {
      this.stopping = true
      await done
    })
    const api = new Api(url.replace(/\/$/, ''), token)
    let ranges = 0
    try {
      while (!this.stopping && (!this.ranges || ranges < this.ranges)) {
        let lease: Lease | null = null
        try {
          lease = await api.lease()
        } catch (error: any) {
          this.logger.warning(`Lease failed: ${error.message}`)
        }
        if (!lease) {
          if (this.ranges) break
          await sleep(IDLE_MS)
          continue
        }
        await this.runRange(api, lease)
        ranges++
      }
    } finally {
      stopped!()
    }
    await this.terminate()
  }

  private async runRange(api: Api, lease: Lease) {
    const { range, rules } = lease
    const started = Date.now()
    const indexer = new RangeIndexer({ group: range.group, start: range.start, end: range.end, rules })
    const totals = { headers: 0, releases: 0, candidates: 0, bytes: 0, postMs: 0 }
    let chunks = 0
    const post = async (batch: EdgeBatch) => {
      const at = Date.now()
      const stored = await api.batch(range.id, batch)
      totals.postMs += Date.now() - at
      totals.bytes += stored.bytes
      totals.headers += batch.headers
      totals.releases += stored.releases
      totals.candidates += stored.candidates
    }
    try {
      for await (const frame of decodeFrames(await api.stream(range.id))) {
        indexer.addText(unzipSync(frame.data).toString('utf8'), frame.to + 1)
        chunks++
        const last = frame.to >= range.end
        if (last) indexer.finish()
        if (last || chunks % this.batchChunks === 0) await post(indexer.take())
        if (this.stopping && !last) {
          // Back to the queue; the next worker resumes at the last batch.
          if (chunks % this.batchChunks !== 0) this.logger.info(`Range ${range.id}: giving it back (${chunks % this.batchChunks} chunks since the last batch are read again).`)
          await api.call(range.id, 'release')
          return
        }
      }
      await api.call(range.id, 'finish')
      const seconds = (Date.now() - started) / 1000
      this.logger.info(
        `${range.group} ${range.start}–${range.end}: ${totals.headers.toLocaleString('en')} headers in ${seconds.toFixed(0)} s ` +
          `(${Math.round(totals.headers / seconds).toLocaleString('en')}/s), ${totals.releases} release chunks, ${totals.candidates} probe candidates; ` +
          `posted ${(totals.bytes / 1e6).toFixed(1)} MB in ${(totals.postMs / 1000).toFixed(1)} s.`
      )
    } catch (error: any) {
      if (error instanceof LeaseLost) {
        this.logger.warning(`Range ${range.id}: ${error.message}`)
        return
      }
      this.logger.error(`Range ${range.id} failed: ${error.message}; giving it back.`)
      await api.call(range.id, 'release').catch(() => {})
    }
  }
}

/** The coordinator's node API, with retries for network trouble and 5xx. */
class Api {
  constructor(
    private base: string,
    private token: string
  ) {}

  private async request(path: string, init: RequestInit = {}): Promise<Response> {
    for (let attempt = 1; ; attempt++) {
      try {
        const response = await fetch(`${this.base}/api/v1/node/${path}`, {
          ...init,
          headers: { 'Authorization': `Bearer ${this.token}`, 'Content-Type': 'application/json', 'Accept': 'application/json', ...init.headers },
        })
        if (response.status === 409) throw new LeaseLost(((await response.json().catch(() => ({}))) as { error?: string }).error ?? 'lease lost')
        if (response.status < 500) return response
        if (attempt >= RETRIES) throw new Error(`${path}: HTTP ${response.status}`)
      } catch (error) {
        if (error instanceof LeaseLost || attempt >= RETRIES) throw error
      }
      await sleep(Math.min(60_000, 1000 * 2 ** attempt))
    }
  }

  async lease(): Promise<Lease | null> {
    const response = await this.request('lease', { method: 'POST', body: '{}' })
    if (response.status === 204) return null
    if (!response.ok) throw new Error(`lease: HTTP ${response.status}`)
    return response.json() as Promise<Lease>
  }

  async stream(rangeId: number): Promise<AsyncIterable<Uint8Array>> {
    const response = await this.request(`ranges/${rangeId}/stream`, { headers: { Accept: 'application/octet-stream' } })
    if (!response.ok || !response.body) throw new Error(`stream: HTTP ${response.status}`)
    return response.body as unknown as AsyncIterable<Uint8Array>
  }

  async batch(rangeId: number, batch: EdgeBatch): Promise<{ releases: number; candidates: number; bytes: number }> {
    // Gzipped: chunk payloads are compressed already, the rest (names, fragments, candidates) compresses well.
    const body = gzipSync(JSON.stringify(toWire(batch)), { level: 1 })
    const response = await this.request(`ranges/${rangeId}/batches`, { method: 'POST', body, headers: { 'Content-Encoding': 'gzip' } })
    if (!response.ok) throw new Error(`batch: HTTP ${response.status}`)
    return { ...((await response.json()) as { releases: number; candidates: number }), bytes: body.length }
  }

  async call(rangeId: number, action: 'finish' | 'release' | 'heartbeat') {
    const response = await this.request(`ranges/${rangeId}/${action}`, { method: 'POST', body: '{}' })
    if (!response.ok) throw new Error(`${action}: HTTP ${response.status}`)
  }
}