///|
/// How a write batch is split before it is sent.
///
/// Chroma refuses a batch larger than the server's `max_batch_size` (it
/// answers 400), and the limit differs per deployment: `GET /pre-flight-checks`
/// reports it when the server implements the endpoint. The client cannot guess
/// it, so it does not guess — a caller who knows the number sets it here, and
/// `write_records` slices the batch into chunks of that size.
///
/// `max_batch_size <= 0` means "send it as one request", which is the default
/// and the behaviour every earlier version had.
pub(all) struct WriteOptions {
  /// Largest number of records per request, when the caller knows the server's
  /// limit. `0` (the default) sends the whole batch at once.
  max_batch_size : Int32
} derive(Eq, Debug)

///|
/// Builds write options; `max_batch_size` defaults to "do not chunk".
pub fn WriteOptions::new(max_batch_size~ : Int32 = 0) -> WriteOptions {
  { max_batch_size, }
}

///|
/// Returns a copy sized for the server's advertised limit.
///
/// Pairs with `preflight`: `WriteOptions::from_preflight(client.preflight())`
/// turns what the server says into what the client sends, and falls back to
/// "do not chunk" when an older server does not advertise a limit.
pub fn WriteOptions::from_preflight(checks : PreflightChecks) -> WriteOptions {
  match checks.max_batch_size {
    Some(size) if size > 0 => { max_batch_size: size, }
    _ => { max_batch_size: 0, }
  }
}

///|
/// Rejects a batch Chroma would refuse for reasons the record itself cannot
/// show, before a request is built.
///
/// `validate_records` checks each record in isolation; these two checks are
/// about the batch as a whole:
///
/// - **duplicate ids inside one request** — Chroma rejects an `add` that names
///   the same id twice, because an insert cannot satisfy "insert twice". The
///   same batch split into two requests works, which is why this is a batch
///   rule and not a per-record one. `upsert` is unaffected: an upsert on an id
///   it just wrote is well defined.
/// - **vector width** — every embedding in a batch must have the length the
///   collection's index was created with. Mismatched widths are a 400 from the
///   server with no hint about which record is wrong, so the first offending
///   record is named here.
pub fn validate_batch(
  records : Array[Record],
  reject_duplicate_ids~ : Bool = true,
) -> Unit raise ChromaError {
  let seen : Map[String, Int32] = Map([])
  let mut expected_dimension = -1
  let mut dimension_owner = ""
  let mut index = 0
  for record in records {
    match seen.get(record.id) {
      Some(first) =>
        if reject_duplicate_ids {
          raise ChromaError::InvalidArgument(
            "record " +
            record.id +
            " appears at index " +
            first.to_string() +
            " and again at index " +
            index.to_string() +
            "; Chroma rejects an add naming the same id twice, send one record per id (or use upsert)",
          )
        }
      None => seen[record.id] = index
    }
    match record.embedding {
      Some(embedding) =>
        if expected_dimension < 0 {
          expected_dimension = embedding.length()
          dimension_owner = record.id
        } else if embedding.length() != expected_dimension {
          raise ChromaError::InvalidArgument(
            "record " +
            record.id +
            " has a " +
            embedding.length().to_string() +
            "-dimensional embedding but record " +
            dimension_owner +
            " has " +
            expected_dimension.to_string() +
            "; a collection's vectors must all share one width",
          )
        }
      None => ()
    }
    index = index + 1
  }
}

///|
/// Splits a batch into the requests `WriteOptions` allows.
///
/// Returns the slices in send order. A batch that fits in one request comes
/// back as a single slice pointing at the original array, so the common case
/// copies nothing.
fn split_batch(
  records : Array[Record],
  options : WriteOptions,
) -> Array[Array[Record]] {
  if options.max_batch_size <= 0 || records.length() <= options.max_batch_size {
    return [records]
  }
  let chunks : Array[Array[Record]] = []
  let mut start = 0
  while start < records.length() {
    let end : Int32 = if start + options.max_batch_size > records.length() {
      records.length()
    } else {
      start + options.max_batch_size
    }
    chunks.push(records[start:end].to_owned())
    start = start + options.max_batch_size
  }
  chunks
}

///|
/// Send one chunk of a write batch: the seam `write_records` hands each slice
/// to, so the sync and async paths share the chunking and differ only here.
pub type ChunkSender = (Array[Record]) -> Unit raise ChromaError

///|
/// The async flavour of `ChunkSender`. Each chunk is awaited before the next
/// is built, so a failure stops the sequence.
pub type AsyncChunkSender = async (Array[Record]) -> Unit raise ChromaError

///|
/// How far a write got before it was interrupted.
///
/// A chunked write can stop in the middle: chunk 3 of 5 is refused and chunks
/// 4 and 5 are never sent, so the caller is left with a partially written
/// collection. This is the value that says so: it carries the error *and* the
/// count of records the server accepted, which is what makes a retry possible
/// — `sent` is a prefix of the batch, so replaying the batch from `sent` under
/// `upsert` is safe.
///
/// `sent` counts accepted chunks only: it is advanced after a chunk's response
/// comes back, never when the chunk is handed to the transport, so a chunk that
/// fails on the way out is *not* counted. Nothing is raised when the whole
/// batch lands, or when it fails before any request is built (validation
/// errors are the caller's arguments, not a partial write).
pub(all) struct PartialFailure {
  /// Records the server accepted before the failure, from `0` to the batch
  /// size. Records `[0, sent)` are stored.
  sent : Int32
  /// What stopped the write: the failing chunk's error.
  cause : ChromaError
}

///|
/// Two failures are the same when they stopped at the same record for the same
/// reason; `ChromaError` compares structurally, so this recurses.
impl Eq for PartialFailure with fn equal(self, other) {
  self.sent == other.sent && self.cause == other.cause
}

///|
/// Turns the error that stopped a chunked write into what the caller sees.
///
/// The first chunk is not a partial write — nothing landed, so the error is
/// returned untouched and callers keep matching on `Api` / `Transport` exactly
/// as they did before chunking existed. From the second chunk on, the batch has
/// a stored prefix and that prefix is the whole point of `PartialFailure`.
fn partial_write(accepted : Int32, error : ChromaError) -> ChromaError {
  if accepted > 0 {
    ChromaError::PartialFailure({ sent: accepted, cause: error, })
  } else {
    error
  }
}

///|
/// The async counterpart of `write_records`, kept separate because MoonBit
/// has no way to await inside a synchronous callback.
///
/// A chunk that raises after earlier chunks landed is re-raised as
/// `ChromaError::PartialFailure`; when the very first chunk fails there is
/// nothing partial about it, so the cause is raised unchanged.
async fn write_records_async(
  records : Array[Record],
  options : WriteOptions,
  reject_duplicate_ids : Bool,
  send_chunk : AsyncChunkSender,
) -> Unit raise ChromaError {
  validate_records(records)
  validate_batch(records, reject_duplicate_ids~)
  let mut accepted = 0
  for chunk in split_batch(records, options) {
    let size = chunk.length()
    send_chunk(chunk) catch {
      error => raise partial_write(accepted, error)
    }
    accepted = accepted + size
  }
}

///|
/// Validates a batch, splits it by `options` and sends each slice.
///
/// Two things are guaranteed on the way out: nothing is sent for a batch that
/// fails validation, and a failure in chunk `n` stops the sequence — the
/// records after it are not written, which is what makes the caller's retry
/// (with the same ids) safe under `upsert`. When `n > 1`, the failure is
/// raised as `PartialFailure` carrying the count of records that did land.
fn write_records(
  records : Array[Record],
  options : WriteOptions,
  reject_duplicate_ids : Bool,
  send_chunk : ChunkSender,
) -> Unit raise ChromaError {
  validate_records(records)
  validate_batch(records, reject_duplicate_ids~)
  let mut accepted = 0
  for chunk in split_batch(records, options) {
    let size = chunk.length()
    send_chunk(chunk) catch {
      error => raise partial_write(accepted, error)
    }
    accepted = accepted + size
  }
}

///|
/// The no-op progress handler, so a call site never branches on "was a
/// handler passed": an optional callback with a default value would make every
/// wrapper wrap it in `Some`.
pub fn no_progress(sent : Int32, total : Int32) -> Unit {
  ignore(sent)
  ignore(total)
}

///|
/// Closes a `where` clause Chroma would misread rather than reject.
///
/// Chroma treats a flat `where` object as a conjunction of its keys
/// (`{"a": .., "b": ..}` is `a AND b`). `{"$and": .., "b": ..}` is not a
/// conjunction Chroma understands: the server answers 400
/// ("Invalid where clause") for most shapes and silently ignores the extra
/// top-level key for some, which is worse because the caller gets records they
/// did not filter for. This folds a top-level key that sits next to `$and` /
/// `$or` into that operator's operand list, so both spellings mean the same
/// thing and nothing is dropped.
///
/// A clause with no operator at the top level is returned unchanged.
pub fn normalize_where(clause : Json) -> Json {
  match clause {
    Json::Object(fields) => {
      let operator_entries : Array[(String, Json)] = []
      let plain_entries : Array[(String, Json)] = []
      for entry in fields {
        let (key, _) = entry
        if key == "$and" || key == "$or" {
          operator_entries.push(entry)
        } else if key == "$nor" {
          // Not a Chroma operator; leave the whole clause alone so the server
          // reports it instead of this function silently dropping it.
          return clause
        } else {
          plain_entries.push(entry)
        }
      }
      if operator_entries.length() != 1 || plain_entries.length() == 0 {
        return clause
      }
      let (operator, operands_json) = operator_entries[0]
      let operands : Array[Json] = match operands_json {
        Json::Array(items) => items.copy()
        _ =>
          // A malformed operand list is the server's to report.
          return clause
      }
      for entry in plain_entries {
        let (key, value) = entry
        let single : Map[String, Json] = Map([])
        single[key] = value
        operands.push(Json::object(single))
      }
      let wrapped : Map[String, Json] = Map([])
      wrapped[operator] = Json::array(operands)
      Json::object(wrapped)
    }
    _ => clause
  }
}

///|
/// Pads a query result so every row has one entry per column.
///
/// A query response is columnar *and* nested: each field is an array with one
/// row per query embedding. For an embedding that matched nothing, Chroma
/// sends `[]` for that row — or omits the field for that row — while the other
/// fields may still carry a row. `QueryResult::from_query_json` reads the
/// first row, and a field that is shorter than `ids` there means "no value for
/// those rows", not "this field was not requested".
///
/// Without the padding, `result.documents[i]` on such a result aborts at
/// runtime (index out of bounds) in a caller that loops over `ids`, which is
/// the natural way to read a result. Here the missing cells become `None`, so
/// the arrays line up with `ids` no matter what the server sent.
pub fn pad_query_result(result : QueryResult) -> QueryResult {
  let rows = result.ids.length()
  {
    ids: result.ids,
    documents: pad_to(result.documents, rows, None),
    embeddings: pad_to(result.embeddings, rows, None),
    distances: pad_to(result.distances, rows, None),
    metadatas: pad_metadata_to(result.metadatas, rows),
  }
}

///|
fn[T] pad_to(values : Array[T?], rows : Int32, filler : T?) -> Array[T?] {
  let out : Array[T?] = []
  for value in values {
    out.push(value)
  }
  while out.length() < rows {
    out.push(filler)
  }
  if out.length() > rows {
    return out[:rows].to_owned()
  }
  out
}

///|
fn pad_metadata_to(values : Array[Metadata], rows : Int32) -> Array[Metadata] {
  let out : Array[Metadata] = []
  for value in values {
    out.push(value)
  }
  while out.length() < rows {
    out.push([])
  }
  if out.length() > rows {
    return out[:rows].to_owned()
  }
  out
}

///|
pub extend WriteOptions with Eq::{not_equal, equal}

///|
pub extend WriteOptions with @debug.Debug::{to_repr}