import type { HttpContext } from '@adonisjs/core/http'
import { Readable } from 'node:stream'
import { setTimeout as sleep } from 'node:timers/promises'
import Leases, { LeaseLost } from '#services/leases'
import BatchStore, { StaleBatch } from '#services/batch_store'
import Spool from '#services/spool'
import { encodeFrame, fromWire, type WireBatch } from '#lib/edge/index'

/** How long the stream waits for a chunk the prefetcher hasn't spooled yet. */
const CHUNK_WAIT_MS = 15 * 60 * 1000

/**
 * The edge workers' API (REBUILD.md, parts 1–2): a lease, the range's
 * chunks as one stream, results per batch, and the end of the range.
 */
export default class NodeController {
  private node(ctx: HttpContext) {
    return Number(ctx.auth.use('node').getUserOrFail().id)
  }

  async lease(ctx: HttpContext) {
    const lease = await Leases.lease(this.node(ctx))
    if (!lease) return ctx.response.noContent()
    return lease
  }

  /** The range's chunks from its resume point, waiting for those still being fetched. */
  async stream(ctx: HttpContext) {
    const nodeId = this.node(ctx)
    const rangeId = Number(ctx.params.id)
    const range = await this.leased(ctx, rangeId, nodeId)
    if (!range) return
    const spans = Spool.chunksOf({ start: Number(range.start_id), end: Number(range.end_id) }, Number(range.next_id))
    async function* frames() {
      for (const span of spans) {
        const waitUntil = Date.now() + CHUNK_WAIT_MS
        while (!(await Spool.has(rangeId, span))) {
          if (Date.now() > waitUntil) throw new Error(`Chunk ${span.from}–${span.to} of range ${rangeId} wasn't spooled in time`)
          await sleep(500)
        }
        yield encodeFrame({ ...span, data: await Spool.read(rangeId, span) })
        // Streaming keeps the lease alive.
        await Leases.heartbeat(rangeId, nodeId)
      }
    }
    ctx.response.header('Content-Type', 'application/octet-stream')
    ctx.response.stream(Readable.from(frames()))
  }

  async batch(ctx: HttpContext) {
    const nodeId = this.node(ctx)
    const rangeId = Number(ctx.params.id)
    try {
      return await BatchStore.store(rangeId, nodeId, fromWire(ctx.request.body() as WireBatch))
    } catch (error) {
      if (error instanceof StaleBatch) return ctx.response.conflict({ error: error.message })
      throw error
    }
  }

  async heartbeat(ctx: HttpContext) {
    return this.lifecycle(ctx, (rangeId, nodeId) => Leases.heartbeat(rangeId, nodeId))
  }

  async finish(ctx: HttpContext) {
    return this.lifecycle(ctx, (rangeId, nodeId) => Leases.finish(rangeId, nodeId))
  }

  async release(ctx: HttpContext) {
    return this.lifecycle(ctx, (rangeId, nodeId) => Leases.release(rangeId, nodeId))
  }

  private async lifecycle(ctx: HttpContext, action: (rangeId: number, nodeId: number) => Promise<void>) {
    try {
      await action(Number(ctx.params.id), this.node(ctx))
      return { ok: true }
    } catch (error) {
      if (error instanceof LeaseLost) return ctx.response.conflict({ error: error.message })
      throw error
    }
  }

  private async leased(ctx: HttpContext, rangeId: number, nodeId: number) {
    try {
      return await Leases.leased(rangeId, nodeId)
    } catch (error) {
      if (error instanceof LeaseLost) {
        ctx.response.conflict({ error: error.message })
        return null
      }
      throw error
    }
  }
}