import { BaseCommand, flags } from '@adonisjs/core/ace'
import type { CommandOptions } from '@adonisjs/core/types/ace'
import { setTimeout as sleep } from 'node:timers/promises'
import { spawn, type ChildProcess } from 'node:child_process'
import { availableParallelism } from 'node:os'
import { promisify } from 'node:util'
import { gzip, unzip } from 'node:zlib'
import { RangeIndexer, decodeFrames, isKeepalive, throttle, toWire, type EdgeBatch, type Lease } from '#lib/edge/index'

const IDLE_MS = 30_000
/** A batch is posted once this many files (completed, or cut at an edge) wait, whatever its articles. */
const MAX_PENDING_FILES = 100_000
const unzipAsync = promisify(unzip)
const gzipAsync = promisify(gzip)
const RETRIES = 8

class LeaseLost extends Error {}
/** A spooled chunk that doesn't inflate (reported; the coordinator fetches it again). */
class CorruptChunk 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-articles
 * articles (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: 'Articles per posted batch (500k: ~4% of releases get a second chunk; LOG.md 2026-10-09)', default: 500_000 })
  declare batchArticles: number

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

  @flags.number({ description: 'Worker processes (one core each; default: EDGE_PROCESSES or 1, 0 = one per core)' })
  declare processes?: number

  @flags.number({ description: 'Download cap for the whole node in Mbit/s, shared by its processes (default: EDGE_MAX_MBIT or none)' })
  declare maxMbit?: 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')
    const requested = this.processes ?? Number(process.env.EDGE_PROCESSES ?? 1)
    const processes = requested === 0 ? availableParallelism() : requested
    if (processes > 1 && !process.env.EDGE_CHILD) return this.supervise(processes)
    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()
  }

  /**
   * One worker per process (a worker is single-threaded): starts `count`
   * children, restarts one that exits unexpectedly, and passes SIGTERM on
   * (each gives its range back).
   */
  private async supervise(count: number) {
    const cap = this.cap()
    this.logger.info(`Starting ${count} worker processes${cap ? ` (download cap ${cap} Mbit/s, ${(cap / count).toFixed(1)} each)` : ''}.`)
    const children = new Set<ChildProcess>()
    let stopping = false
    let stopped: () => void
    const done = new Promise<void>((resolve) => (stopped = resolve))
    const start = (slot: number) => {
      const args = [
        process.argv[1],
        'edge:node',
        `--batch-articles=${this.batchArticles}`,
        ...(this.ranges ? [`--ranges=${this.ranges}`] : []),
        // Each child gets its share of the node's cap.
        ...(cap ? [`--max-mbit=${cap / count}`] : []),
      ]
      const child = spawn(process.execPath, [...process.execArgv, ...args], { env: { ...process.env, EDGE_CHILD: String(slot) }, stdio: 'inherit' })
      children.add(child)
      child.on('exit', (code) => {
        children.delete(child)
        if (stopping || (this.ranges && code === 0)) {
          if (children.size === 0) stopped()
          return
        }
        this.logger.warning(`Worker ${slot} exited (${code}); starting it again in 10 s.`)
        setTimeout(() => !stopping && start(slot), 10_000)
      })
    }
    for (let slot = 1; slot <= count; slot++) start(slot)
    this.app.terminating(async () => {
      stopping = true
      for (const child of children) child.kill('SIGTERM')
      await done
    })
    await done
    await this.terminate()
  }

  /** The node's download cap in Mbit/s, or null. */
  private cap(): number | null {
    const mbit = this.maxMbit ?? (process.env.EDGE_MAX_MBIT ? Number(process.env.EDGE_MAX_MBIT) : null)
    return mbit && mbit > 0 ? mbit : null
  }

  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, resumeAt: range.next, rules })
    const totals = { headers: 0, releases: 0, candidates: 0, bytes: 0, postMs: 0 }
    let batchFrom = range.next
    let reachedEnd = false
    // Posts go in order, but the worker indexes on meanwhile (a post took 1–16 s; LOG.md 2026-10-09).
    let posting: Promise<void> = Promise.resolve()
    const post = (batch: EdgeBatch) => {
      const previous = posting
      posting = (async () => {
        await previous
        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
      })()
      // A failed post stops the range at the next chunk (a lost lease: nothing more will be stored).
      posting.catch((error) => (postError ??= error))
    }
    let postError: Error | null = null
    const handle = async (chunk: { from: number; to: number; text: Promise<Buffer> }) => {
      if (postError) throw postError
      let text: Buffer
      try {
        text = await chunk.text
      } catch (error: any) {
        await posting.catch(() => {})
        await api.corrupt(range.id, chunk.from)
        throw new CorruptChunk(`chunk ${chunk.from}–${chunk.to} doesn't inflate (${error.message}); reported, the range goes back`)
      }
      indexer.addText(text.toString('utf8'), chunk.to + 1)
      reachedEnd = chunk.to >= range.end
      if (reachedEnd) indexer.finish()
      // Also early when many files wait: a picture group's 500k articles made batches over the server's 256 MB limit (HTTP 413).
      if (reachedEnd || chunk.to + 1 - batchFrom >= this.batchArticles || indexer.pending >= MAX_PENDING_FILES) {
        post(indexer.take())
        batchFrom = chunk.to + 1
      }
    }
    try {
      // Inflating runs on a thread-pool thread: the next chunk inflates while this one is indexed.
      let pending: { from: number; to: number; text: Promise<Buffer> } | null = null
      const cap = this.cap()
      const stream = await api.stream(range.id)
      for await (const frame of decodeFrames(cap ? throttle(stream, (cap * 1e6) / 8) : stream)) {
        if (this.stopping) break
        if (isKeepalive(frame)) continue
        const next = { from: frame.from, to: frame.to, text: unzipAsync(frame.data) }
        next.text.catch(() => {})
        if (pending) await handle(pending)
        pending = next
      }
      if (pending && !this.stopping) await handle(pending)
      await posting
      if (!reachedEnd) {
        // Stopping: back to the queue; the next worker resumes at the last stored batch.
        this.logger.info(`Range ${range.id}: giving it back.`)
        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) {
      await posting.catch(() => {})
      if (error instanceof LeaseLost || error instanceof CorruptChunk) {
        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 (on a thread-pool thread): chunk payloads are compressed already, the rest compresses well.
    const body = await gzipAsync(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 corrupt(rangeId: number, from: number) {
    const response = await this.request(`ranges/${rangeId}/corrupt`, { method: 'POST', body: JSON.stringify({ from }) })
    if (!response.ok) throw new Error(`corrupt: HTTP ${response.status}`)
  }

  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}`)
  }
}