// External sort for ORDER BY.
//
// Sorted runs: every source batch is sorted in memory and split into
// chunks; chunks accumulate until ~budget/4 bytes, then flush as one
// run file. Reduce: pairwise merges until at most FANIN runs remain
// (memory per pass = two runs). Final pass: k-way merge streams rows
// through projection, DISTINCT and the OFFSET/LIMIT window, so the
// result — with a LIMIT — never materializes beyond the output window;
// without LIMIT the full result is the query's answer and is
// materialized as such.
//
// The budget bounds the run state, not the base tables (scan laziness
// is out of the sort's scope) and not the query result itself.
//
// Nothing here reads a batch array any more. `build_runs` takes a source and
// fills runs as batches arrive, so the input is never an array; `MergeCursor`
// is the final k-way merge as a cursor, which a caller that can hand rows on
// (a `ResultStream`) drives one batch at a time, and which `final_merge` drains
// for the collected answer. Both started as the same code: an array of kept
// batches and one `final_merge` returning one batch.

///|
const FANIN : Int32 = 8

///|
priv struct RunReader {
  batches : Array[@types.DataBatch]
  mut bi : Int32
  mut ri : Int32
  mut row : Array[@types.Scalar] // current pre-projection row, empty when exhausted
  mut keys : Array[@types.Scalar] // key values of the current row
  width : Int32
  types : Array[@types.DataType]
  keyspecs : Array[(PhysExpr, Bool, Bool?)]
}

///|
fn RunReader::exhausted(self : RunReader) -> Bool {
  self.row.length() == 0
}

///|
fn RunReader::advance(self : RunReader) -> Unit raise @types.SqlError {
  self.ri += 1
  if self.ri >= self.batches[self.bi].row_count() {
    self.bi += 1
    self.ri = 0
    if self.bi >= self.batches.length() {
      self.row = []
      self.keys = []
      return
    }
  }
  let batch = self.batches[self.bi]
  let row : Array[@types.Scalar] = []
  for c in 0.. @types.Scalar = fn(i) { row[i] }
  let keys : Array[@types.Scalar] = []
  for kv in self.keyspecs {
    keys.push(eval_with(kv.0, get))
  }
  self.keys = keys
}

///|
fn RunReader::make(
  batches : Array[@types.DataBatch],
  keyspecs : Array[(PhysExpr, Bool, Bool?)],
) -> RunReader raise @types.SqlError {
  let width = if batches.length() > 0 {
    batches[0].columns().length()
  } else {
    0
  }
  let types : Array[@types.DataType] = []
  if batches.length() > 0 {
    for col in batches[0].columns() {
      types.push(col.dtype())
    }
  }
  let r : RunReader = {
    batches,
    bi: 0,
    ri: -1,
    row: [],
    keys: [],
    width,
    types,
    keyspecs,
  }
  if batches.length() > 0 {
    r.advance()
  }
  r
}

///|
/// Compare two key rows under (desc, nulls_first) specs; ties return 0
/// (caller breaks them, e.g. by run order for global stability).
fn cmp_key_rows(
  a : Array[@types.Scalar],
  b : Array[@types.Scalar],
  keyspecs : Array[(PhysExpr, Bool, Bool?)],
) -> Int32 {
  for k, kv in keyspecs {
    let x = a[k]
    let y = b[k]
    let nf = match kv.2 {
      Some(f) => f
      None => false
    }
    let ord : Int32 = match (x, y) {
      (@types.Null, @types.Null) => 0
      (@types.Null, _) => if nf { -1 } else { 1 }
      (_, @types.Null) => if nf { 1 } else { -1 }
      _ => {
        let c = @types.Scalar::compare_nn(x, y)
        if kv.1 {
          -c
        } else {
          c
        }
      }
    }
    if ord != 0 {
      return ord
    }
  }
  0
}

///|
fn chunk_of(
  batch : @types.DataBatch,
  from : Int32,
  until : Int32,
) -> @types.DataBatch {
  let indices : Array[Int32] = []
  for r in from.. @types.DataBatch? raise @types.SqlError,
  keys : Array[(PhysExpr, Bool, Bool?)],
  budget : Int32,
  store : SpillStore,
) -> RunBuild raise @types.SqlError {
  let writer : RunWriter = {
    store,
    budget,
    sort_keys: keys,
    acc: [],
    acc_bytes: 0,
    ids: [],
  }
  while next() is Some(batch) {
    if keys.length() > 0 && batch.row_count() > 1 {
      let sorted = sort_exprs(batch, keys, None, None)
      // split oversized sorted batches so one run cannot balloon past
      // the accumulator by itself
      let bytes = estimate_batch_bytes(sorted)
      let per_row = if sorted.row_count() > 0 {
        bytes / sorted.row_count()
      } else {
        1
      }
      let chunk_rows = if per_row > 0 {
        budget / 8 / per_row
      } else {
        sorted.row_count()
      }
      if chunk_rows >= sorted.row_count() || chunk_rows < 1 {
        writer.push(sorted)
      } else {
        let mut from = 0
        while from < sorted.row_count() {
          let until = if from + chunk_rows < sorted.row_count() {
            from + chunk_rows
          } else {
            sorted.row_count()
          }
          writer.push(chunk_of(sorted, from, until))
          from = until
        }
      }
    } else {
      writer.push(batch)
    }
  }
  if writer.ids.length() == 0 {
    return { ids: [], acc: writer.acc, }
  }
  writer.flush()
  { ids: writer.ids, acc: [], }
}

///|
/// External sort driving the whole ORDER BY tail: returns the final
/// projected (and DISTINCT/OFFSET/LIMIT-filtered) batch. Falls back to
/// the in-memory pipeline when everything fits within the run
/// accumulator.
///
/// The input is a source, not an array: `next` is called until it returns
/// `None`, and the batches go into runs as they come. `pre_types` is the batch
/// layout of that input (the ORDER BY keys and the projection both read it), so
/// an input of no batches still yields a typed empty answer.
fn external_sort_project(
  next : () -> @types.DataBatch? raise @types.SqlError,
  pre_types : Array[@types.DataType],
  keys : Array[(PhysExpr, Bool, Bool?)],
  out_exprs : Array[PhysExpr],
  out_types : Array[@types.DataType],
  distinct : Bool,
  offset : Int32?,
  limit : Int32?,
  budget : Int32,
  store : SpillStore,
  stats : Array[(String, Int64)],
) -> @types.DataBatch raise @types.SqlError {
  let built = build_runs(next, keys, budget, store)
  // everything fit: no run was flushed — take the in-memory path
  if built.ids.length() == 0 {
    let pre = concat_batches(built.acc, pre_types)
    let sorted = sort_exprs(pre, keys, None, None)
    return finish_sorted(sorted, out_exprs, out_types, distinct, offset, limit)
  }
  if stats.length() > 0 {
    stats.push(("spill runs", built.ids.length().to_int64()))
  }
  let owned : Array[Int32] = []
  for id in built.ids {
    owned.push(id)
  }
  errdefer remove_ids(store, owned)
  let inputs = reduce_runs(built.ids, keys, store, owned)
  let out = final_merge(
    inputs, keys, out_exprs, out_types, distinct, offset, limit, store,
  )
  // every file this sort wrote is dead now: the merge inputs have been read
  // whole and the intermediates were consumed by the passes that made them
  // (removal is best-effort, so removing both sets is safe)
  remove_ids(store, owned)
  out
}

///|
/// Pairwise-merge runs down to FANIN, and answer the files the final k-way
/// merge must read. Every file this writes is appended to `owned` — the
/// caller's list of what still has to be removed — because the intermediate
/// passes delete the pair they consume (`merge_runs` does) and only the last
/// level survives to be merged.
///
/// Memory per pass is two runs, so the reduction is bounded by the budget even
/// though it is not lazy: a k-way merge needs at most FANIN cursors, and a
/// cursor holds one run.
fn reduce_runs(
  runs0 : Array[Int32],
  keys : Array[(PhysExpr, Bool, Bool?)],
  store : SpillStore,
  owned : Array[Int32],
) -> Array[Int32] raise @types.SqlError {
  let mut cur = runs0
  while cur.length() > FANIN {
    let next : Array[Int32] = []
    let mut i = 0
    while i < cur.length() {
      if i + 1 < cur.length() {
        let id = merge_runs(cur[i], cur[i + 1], keys, store)
        owned.push(id)
        next.push(id)
      } else {
        next.push(cur[i])
      }
      i += 2
    }
    cur = next
  }
  cur
}

///|
/// The final k-way merge as a cursor: one batch of the ordered result per
/// `next`, `None` at the end. Projection, DISTINCT and the OFFSET/LIMIT window
/// are applied per row, so a windowed merge stops as soon as the window is full
/// and the caller never holds more than one output batch.
///
/// This is the whole of `final_merge` except the drain, which is why the
/// collected answer and a streamed one cannot disagree about a row: they run
/// the same cursor and one of them concatenates.
priv struct MergeCursor {
  readers : Array[RunReader]
  keys : Array[(PhysExpr, Bool, Bool?)]
  out_exprs : Array[PhysExpr]
  out_types : Array[@types.DataType]
  distinct : Bool
  offset : Int32?
  limit : Int32?
  seen : Map[String, Bool]
  mut skipped : Int32
  mut taken : Int32
  mut builder : OutputBuilder
  mut finished : Bool
}

///|
fn merge_open(
  runs : Array[Int32],
  keys : Array[(PhysExpr, Bool, Bool?)],
  out_exprs : Array[PhysExpr],
  out_types : Array[@types.DataType],
  distinct : Bool,
  offset : Int32?,
  limit : Int32?,
  store : SpillStore,
) -> MergeCursor raise @types.SqlError {
  let readers : Array[RunReader] = []
  for id in runs {
    readers.push(
      RunReader::make(
        @types.DataBatch::from_spill_bytes((store.read)(id)),
        keys,
      ),
    )
  }
  {
    readers,
    keys,
    out_exprs,
    out_types,
    distinct,
    offset,
    limit,
    seen: Map([]),
    skipped: 0,
    taken: 0,
    builder: OutputBuilder::new(out_types, 4096),
    finished: false,
  }
}

///|
fn MergeCursor::window_full(self : MergeCursor) -> Bool {
  match self.limit {
    Some(l) => self.taken >= l
    None => false
  }
}

///|
fn MergeCursor::next(
  self : MergeCursor,
) -> @types.DataBatch? raise @types.SqlError {
  if self.finished {
    return None
  }
  while !self.window_full() {
    // pick the smallest current row; ties prefer the lower run index
    let mut best = -1
    for r, rd in self.readers {
      if rd.exhausted() {
        continue
      }
      if best == -1 ||
        cmp_key_rows(rd.keys, self.readers[best].keys, self.keys) < 0 {
        best = r
      }
    }
    if best == -1 {
      break
    }
    let row = self.readers[best].row
    if self.offset is Some(o) && self.skipped < o {
      self.skipped += 1
    } else {
      let get : (Int32) -> @types.Scalar = fn(i) { row[i] }
      let projected : Array[@types.Scalar] = []
      for e in self.out_exprs {
        projected.push(eval_with(e, get))
      }
      let mut keep = true
      if self.distinct {
        let kv = canonical_key(projected)
        if self.seen.get(kv) is Some(_) {
          keep = false
        } else {
          self.seen[kv] = true
        }
      }
      if keep {
        self.builder.push_row(projected)
        self.taken += 1
        if self.builder.full() {
          let batch = self.builder.build()
          self.builder = OutputBuilder::new(self.out_types, 4096)
          self.readers[best].advance()
          return Some(batch)
        }
      }
    }
    self.readers[best].advance()
  }
  self.finished = true
  if self.builder.empty() {
    None
  } else {
    let batch = self.builder.build()
    self.builder = OutputBuilder::new(self.out_types, 4096)
    Some(batch)
  }
}

///|
/// Drain the merge into one batch — the collected answer.
fn final_merge(
  runs : Array[Int32],
  keys : Array[(PhysExpr, Bool, Bool?)],
  out_exprs : Array[PhysExpr],
  out_types : Array[@types.DataType],
  distinct : Bool,
  offset : Int32?,
  limit : Int32?,
  store : SpillStore,
) -> @types.DataBatch raise @types.SqlError {
  let cursor = merge_open(
    runs, keys, out_exprs, out_types, distinct, offset, limit, store,
  )
  let chunks : Array[@types.DataBatch] = []
  while cursor.next() is Some(batch) {
    chunks.push(batch)
  }
  concat_batches(chunks, out_types)
}

///|
/// Merge two run files into a new one (row payloads only, ordering
/// preserved). Memory = the two inputs plus the merged buffer.
fn merge_runs(
  a : Int32,
  b : Int32,
  keys : Array[(PhysExpr, Bool, Bool?)],
  store : SpillStore,
) -> Int32 raise @types.SqlError {
  let ra = RunReader::make(
    @types.DataBatch::from_spill_bytes((store.read)(a)),
    keys,
  )
  let rb = RunReader::make(
    @types.DataBatch::from_spill_bytes((store.read)(b)),
    keys,
  )
  let chunks : Array[@types.DataBatch] = []
  let mut builder = OutputBuilder::new(ra.types, 4096)
  while !(ra.exhausted() && rb.exhausted()) {
    let take_a = if ra.exhausted() {
      false
    } else if rb.exhausted() {
      true
    } else {
      cmp_key_rows(ra.keys, rb.keys, keys) <= 0 // ties prefer the earlier run: global stability
    }
    let row = if take_a { ra.row } else { rb.row }
    builder.push_row(row)
    if builder.full() {
      chunks.push(builder.build())
      builder = OutputBuilder::new(ra.types, 4096)
    }
    if take_a {
      ra.advance()
    } else {
      rb.advance()
    }
  }
  if !builder.empty() {
    chunks.push(builder.build())
  }
  let id = (store.write)(@types.DataBatch::to_spill_bytes(chunks))
  (store.remove)(a)
  (store.remove)(b)
  id
}

///|
/// Column builder for streamed output rows.
priv struct OutputBuilder {
  types : Array[@types.DataType]
  cap : Int32
  mut filled : Int32
  cols : Array[@types.ColumnVector]
}

///|
fn OutputBuilder::new(
  types : Array[@types.DataType],
  cap : Int32,
) -> OutputBuilder {
  let cols : Array[@types.ColumnVector] = []
  for dt in types {
    cols.push(@types.ColumnVector::make(dt, cap))
  }
  { types, cap, filled: 0, cols, }
}

///|
fn OutputBuilder::full(self : OutputBuilder) -> Bool {
  self.filled == self.cap
}

///|
fn OutputBuilder::empty(self : OutputBuilder) -> Bool {
  self.filled == 0
}

///|
fn OutputBuilder::push_row(
  self : OutputBuilder,
  row : Array[@types.Scalar],
) -> Unit {
  for c in 0.. @types.DataBatch {
  @types.DataBatch::new(self.cols, self.filled)
}

///|
/// The in-memory tail for the everything-fit fallback: sort, project,
/// dedupe, window.
fn finish_sorted(
  sorted : @types.DataBatch,
  out_exprs : Array[PhysExpr],
  out_types : Array[@types.DataType],
  distinct : Bool,
  offset : Int32?,
  limit : Int32?,
) -> @types.DataBatch raise @types.SqlError {
  let projected = project(sorted, out_exprs, out_types)
  let projected = if distinct { dedupe_rows(projected) } else { projected }
  order_limit(projected, [], limit, offset)
}