///|
/// A record stored in a collection.
///
/// `embedding` is `None` for a record that carries only a document. Chroma
/// needs a vector for every record, so a batch with a missing embedding is
/// rejected locally instead of by the server (see `validate_records`).
pub(all) struct Record {
  id : String
  document : String?
  embedding : Embedding?
  metadata : Metadata
} derive(Eq, Debug)

///|
/// Builds a record that carries a document and no embedding.
///
/// Chroma 1.x has no server-side embedding, so a record built this way can
/// only be stored once an `embedding` is attached via `Record::with_embedding`.
pub fn Record::of_document(
  id : String,
  document : String,
  metadata~ : Metadata = [],
) -> Record {
  { id, document: Some(document), embedding: None, metadata, }
}

///|
/// Builds a record with a caller-supplied embedding.
pub fn Record::of_embedding(
  id : String,
  embedding : Embedding,
  metadata~ : Metadata = [],
) -> Record {
  { id, document: None, embedding: Some(embedding), metadata, }
}

///|
/// Returns a copy of the record with `embedding` attached.
///
/// Use it to add a vector to a record built by `Record::of_document`; a
/// record with explicit coordinates is built with `Record::of_embedding`.
pub fn Record::with_embedding(self : Record, embedding : Embedding) -> Record {
  { ..self, embedding: Some(embedding), }
}

///|
/// Returns a copy of the record carrying `document` as well.
///
/// Pairs with `Record::of_embedding` when the vector is known first and the
/// text it stands for second.
pub fn Record::with_document(self : Record, document : String) -> Record {
  { ..self, document: Some(document), }
}

///|
/// Returns a copy of the record carrying `metadata` instead of its own.
pub fn Record::with_metadata(self : Record, metadata : Metadata) -> Record {
  { ..self, metadata, }
}

///|
/// Serializes a batch of records into the columnar body Chroma expects.
///
/// Chroma does not take records one by one: the add/upsert/update models are
/// parallel arrays (`ids`, `embeddings`, `documents`, `metadatas`) that must
/// be the same length and aligned by index. A column is emitted only when it
/// carries data, and a `null` marks a record that has no value for it — that
/// position-preserving `null` is what keeps the arrays aligned.
///
/// `embeddings` is the exception: Chroma requires the field and rejects a
/// batch where it is missing, null or empty, so a record without coordinates
/// has no representation here. `validate_records` refuses such a batch
/// before it is serialized.
pub fn records_to_columns(records : Array[Record]) -> Json {
  let ids : Array[Json] = []
  let embeddings : Array[Json] = []
  let documents : Array[Json] = []
  let metadatas : Array[Json] = []
  let mut has_documents = false
  let mut has_metadatas = false
  for record in records {
    ids.push(Json::string(record.id))
    embeddings.push(
      match record.embedding {
        Some(embedding) => embedding_to_json(embedding)
        None => Json::null()
      },
    )
    match record.document {
      Some(document) => {
        has_documents = true
        documents.push(Json::string(document))
      }
      None => documents.push(Json::null())
    }
    if record.metadata.length() > 0 {
      has_metadatas = true
      metadatas.push(metadata_to_json(record.metadata))
    } else {
      metadatas.push(Json::null())
    }
  }
  let body : Map[String, Json] = Map([])
  body["ids"] = Json::array(ids)
  body["embeddings"] = Json::array(embeddings)
  if has_documents {
    body["documents"] = Json::array(documents)
  }
  if has_metadatas {
    body["metadatas"] = Json::array(metadatas)
  }
  Json::object(body)
}

///|
/// A row returned by `get` or `query`.
pub(all) struct QueryResult {
  ids : Array[String]
  documents : Array[String?]
  embeddings : Array[Embedding?]
  metadatas : Array[Metadata]
  distances : Array[Double?]
} derive(Eq, Debug)

///|
/// Where-clause filters for `get`, `update` and `delete`.
///
/// These mirror Chroma's `where` / `where_document` JSON filters.
pub(all) struct Filter {
  where_clause : Json?
  where_document_clause : Json?
  ids : Array[String]
  limit : Int32?
  offset : Int32?
} derive(Eq, Debug)

///|
/// Builds an empty filter that matches everything.
pub fn Filter::new(
  where_clause~ : Json? = None,
  where_document_clause~ : Json? = None,
  ids~ : Array[String] = [],
  limit~ : Int32? = None,
  offset~ : Int32? = None,
) -> Filter {
  { where_clause, where_document_clause, ids, limit, offset, }
}

///|
/// Adds a metadata equality condition to the filter.
///
/// Repeated calls accumulate instead of replacing each other, the way
/// chroma-go's `WithWhere(EqString(k, v))` options append to the same
/// `where` clause: `with_metadata("a", ..).with_metadata("b", ..)` matches
/// records where both hold, not just the last one.
///
/// The second call wraps both tests in `$and`, because Chroma reads a flat
/// object's keys as alternatives and rejects more than one of them outright
/// ("Invalid where clause") rather than intersecting them. Calling it twice
/// with the *same* key works too: the two tests end up as separate conjuncts
/// and nothing is overwritten.
pub fn Filter::with_metadata(
  self : Filter,
  key : String,
  value : MetadataValue,
) -> Filter {
  let clause = match self.where_clause {
    // A clause this method built itself is already a conjunction: appending
    // to its operand list keeps the shape flat instead of nesting `$and`s.
    Some(Json::Object(fields)) if fields.contains("$and") => {
      let operands = match fields.get("$and") {
        Some(Json::Array(items)) => items
        _ => []
      }
      operands.push(where_eq(key, value))
      where_and(operands)
    }
    Some(previous) => where_and([previous, where_eq(key, value)])
    None => where_eq(key, value)
  }
  { ..self, where_clause: Some(clause), }
}

///|
/// Adds a limit to the filter.
pub fn Filter::with_limit(self : Filter, limit : Int32) -> Filter {
  { ..self, limit: Some(limit), }
}

///|
/// Serializes `get`'s body: the filter fields plus the `include` array.
///
/// This is deliberately not what `update` and `delete` send. Neither endpoint
/// takes an `include` — it only decides which *columns come back*, and a write
/// returns nothing — so putting it in their body is at best ignored and at
/// worst a 400. They build their own bodies through `filter_clauses`, which
/// drops `include` (and `offset`, which no write endpoint accepts either).
///
/// A write body is built from `filter_clauses` directly rather than from this
/// function: `where` clauses are the part all three share, `include` is not
/// (see `docs/protocol-constraints.md` §4).
fn Filter::to_json(self : Filter, include_fields : Include) -> Json {
  let body = match self.filter_clauses() {
    Json::Object(fields) => fields.copy()
    _ => Map([])
  }
  body["include"] = include_fields.to_json()
  Json::object(body)
}

///|
/// Serializes what `get`, `update` and `delete` send in common: the
/// normalized `where` / `where_document` clauses, `ids`, `limit` and
/// `offset`.
///
/// Writing `where` in exactly one place is the point: every clause the client
/// sends goes through `normalize_where` here, so there is no call site that
/// can spell one differently (see `docs/protocol-constraints.md` §4).
///
/// The two write endpoints do not take the paging fields either: `get`
/// validates `limit` / `offset` before building a body, and `update` /
/// `delete` drop whatever paging window the caller left on the filter.
fn Filter::filter_clauses(self : Filter) -> Json {
  let body : Map[String, Json] = Map([])
  match self.where_clause {
    // Normalized here rather than at each call site: `get`, `update` and
    // `delete` all go through this function, so a clause can only be spelled
    // one way on the wire (see `normalize_where`).
    Some(value) => body["where"] = normalize_where(value)
    None => ()
  }
  match self.where_document_clause {
    Some(value) => body["where_document"] = normalize_where(value)
    None => ()
  }
  if self.ids.length() > 0 {
    let ids : Array[Json] = []
    for id in self.ids {
      ids.push(Json::string(id))
    }
    body["ids"] = Json::array(ids)
  }
  match self.limit {
    Some(value) => body["limit"] = Json::number(value.to_double())
    None => ()
  }
  match self.offset {
    Some(value) => body["offset"] = Json::number(value.to_double())
    None => ()
  }
  Json::object(body)
}

///|
/// Adds records to a collection.
///
/// Every record must carry either a `document` or an `embedding`, otherwise
/// Chroma rejects the request.
/// `options.max_batch_size` splits the batch into several requests; a batch
/// that fails `validate_batch` is not sent at all. `on_chunk` reports how many
/// records have landed after each request, so a failure in chunk `n` tells the
/// caller exactly which prefix of the batch is already stored.
pub fn[T : Transport] ChromaClient::add(
  self : ChromaClient[T],
  collection_id : String,
  records : Array[Record],
  options~ : WriteOptions = WriteOptions::new(),
  on_chunk~ : ChunkProgress = no_progress,
) -> Unit raise ChromaError {
  self.write_chunks(
    collection_id,
    records,
    "add",
    options,
    reject_duplicate_ids=true,
    on_chunk~,
  )
}

///|
/// Posts one write chunk, the sync half of `write_records`.
fn[T : Transport] ChromaClient::write_chunks(
  self : ChromaClient[T],
  collection_id : String,
  records : Array[Record],
  endpoint : String,
  options : WriteOptions,
  reject_duplicate_ids~ : Bool,
  on_chunk~ : ChunkProgress = no_progress,
) -> Unit raise ChromaError {
  let path = self.collection_path(collection_id) + "/" + endpoint
  // `sent` lives outside the closure and is only advanced after a chunk has
  // been accepted, so a raised error leaves it at "what actually landed".
  let sent : Ref[Int32] = Ref(0)
  let total = records.length()
  write_records(records, options, reject_duplicate_ids, chunk => {
    ignore(self.send_json(Post, path, body=records_to_columns(chunk)))
    sent.val = sent.val + chunk.length()
    on_chunk(sent.val, total)
  })
}

///|
/// Inserts or replaces records by id.
/// An upsert may name the same id twice: the second write is defined by the
/// id already existing, so only the vector-width rule of `validate_batch`
/// applies.
pub fn[T : Transport] ChromaClient::upsert(
  self : ChromaClient[T],
  collection_id : String,
  records : Array[Record],
  options~ : WriteOptions = WriteOptions::new(),
  on_chunk~ : ChunkProgress = no_progress,
) -> Unit raise ChromaError {
  self.write_chunks(
    collection_id,
    records,
    "upsert",
    options,
    reject_duplicate_ids=false,
    on_chunk~,
  )
}

///|
/// Updates the document, metadata or embedding of existing records.
///
/// Records are matched by id; ids that do not exist are ignored by Chroma.
pub fn[T : Transport] ChromaClient::update(
  self : ChromaClient[T],
  collection_id : String,
  records : Array[Record],
  filter~ : Filter = Filter::new(),
) -> Unit raise ChromaError {
  // An update patches the records it names, so a missing embedding or
  // document is the normal case rather than an error. Chroma's update model
  // accepts `null` in a column to mean "leave this field alone", which is
  // exactly what `records_to_columns` emits.
  validate_records(records, require_embedding=false)
  validate_batch(records, reject_duplicate_ids=false)
  // An update is a write: it carries the records plus the clauses that pick
  // which ones, and nothing that describes a response. That rules out
  // `include` (it selects *returned* columns) and `offset` (no write endpoint
  // pages), both of which `filter_clauses` holds back.
  let body : Map[String, Json] = match records_to_columns(records) {
    Json::Object(fields) => fields.copy()
    _ => Map([])
  }
  for entry in update_clauses(filter) {
    let (key, value) = entry
    body[key] = value
  }
  let path = self.collection_path(collection_id) + "/update"
  ignore(self.send_json(Post, path, body=Json::object(body)))
}

///|
/// Deletes the records a filter names.
///
/// Like chroma-go's `Collection.Delete`, at least one of `ids`, `where` or
/// `where_document` is required: an empty filter is refused locally instead of
/// silently emptying the collection. `limit` is forwarded, and only together
/// with a `where` / `where_document` clause, which is the one case Chroma
/// accepts it in.
pub fn[T : Transport] ChromaClient::delete(
  self : ChromaClient[T],
  collection_id : String,
  filter~ : Filter = Filter::new(),
) -> Unit raise ChromaError {
  let body = delete_body(filter)
  let path = self.collection_path(collection_id) + "/delete"
  ignore(self.send_json(Post, path, body~))
}

///|
/// The `where` / `where_document` clauses of an `update`.
///
/// `update` takes no `ids` / `include` / `limit` / `offset`: it patches the
/// records it is given, and those are named by the `id` column of the body
/// plus an optional clause. Dropping the rest is what keeps a body the server
/// accepts from carrying a field it would reject.
fn update_clauses(filter : Filter) -> Array[(String, Json)] {
  let clauses : Array[(String, Json)] = []
  match filter.where_clause {
    Some(value) => clauses.push(("where", normalize_where(value)))
    None => ()
  }
  match filter.where_document_clause {
    Some(value) => clauses.push(("where_document", normalize_where(value)))
    None => ()
  }
  clauses
}

///|
/// Builds a delete body, refusing a filter that would delete everything.
fn delete_body(filter : Filter) -> Json raise ChromaError {
  if filter.ids.length() == 0 {
    if filter.where_clause is None {
      if filter.where_document_clause is None {
        raise ChromaError::InvalidArgument(
          "delete needs at least one filter: pass ids, a where clause or a where_document clause; " +
          "an empty filter would delete every record in the collection",
        )
      }
    }
  }
  let has_clause = match (filter.where_clause, filter.where_document_clause) {
    (None, None) => false
    _ => true
  }
  let limit = match filter.limit {
    Some(value) => {
      if value <= 0 {
        raise ChromaError::InvalidArgument(
          "delete limit must be greater than zero, got " + value.to_string(),
        )
      }
      if !has_clause {
        raise ChromaError::InvalidArgument(
          "delete limit can only be used together with a where or where_document clause",
        )
      }
      value
    }
    None => 0
  }
  // A delete names records and asks for nothing back, so its body has the
  // clauses and never an `include` / `offset` (see `update`).
  let body = match filter.filter_clauses() {
    Json::Object(fields) => fields.copy()
    _ => Map([])
  }
  // An offset would page the delete, which no write endpoint does: it is
  // dropped rather than forwarded.
  ignore(body.remove("offset"))
  // `limit` is re-stamped from the validated value, so a filter that names
  // records by `ids` only (where Chroma refuses a `limit`) cannot leak one.
  if limit > 0 {
    body["limit"] = Json::number(limit.to_double())
  }
  Json::object(body)
}

///|
/// Fetches records by filter, or the whole collection when the filter is
/// empty.
///
/// Mirrors chroma-go's `Collection.Get`: `include_fields` decides which
/// columns come back and defaults to `Include::get_default()` (documents and
/// metadata). It is deliberately not `Include::default()`, the query default:
/// that one asks for `distances`, and Chroma rejects `POST /get` carrying them
/// (`distances are not supported in get`). `limit` /
/// `offset` page a result set. The two are independent — chroma-go's own
/// example pages with `WithLimit(100), WithOffset(200)`, an offset past the
/// page size — so neither is clamped against the other.
pub fn[T : Transport] ChromaClient::get(
  self : ChromaClient[T],
  collection_id : String,
  filter~ : Filter = Filter::new(),
  include_fields~ : Include = Include::get_default(),
) -> QueryResult raise ChromaError {
  validate_paging(filter.limit, filter.offset)
  let body = filter.to_json(include_fields)
  let path = self.collection_path(collection_id) + "/get"
  QueryResult::from_json(self.send_json(Post, path, body~))
}

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

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

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

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

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

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

///|
/// Rejects record batches Chroma would refuse, before a request is built.
///
/// Chroma answers 400 for an empty batch and for a blank id, and its
/// add/upsert models make `embeddings` mandatory, so a batch whose records
/// have no vectors is refused too. All of these are local mistakes with local
/// fixes, so they are caught here instead of after a round-trip.
///
/// `require_embedding` is `false` for `update`, where omitting a field is how
/// a caller says "change only what I passed".
pub fn validate_records(
  records : Array[Record],
  require_embedding~ : Bool = true,
) -> Unit raise ChromaError {
  validate_records_per_record(records, require_embedding~)
}

///|
/// Reports progress of a chunked write, once per request that succeeded.
///
/// A caller that writes 5,000 records in chunks of 500 has no way to tell how
/// far a failure got unless the client says so: `on_chunk` is called after
/// each slice is accepted by the server, with the number of records written so
/// far and the batch's total. When a later chunk fails, the count is what has
/// already landed and the ids of the failed chunk are the ones to retry.
pub type ChunkProgress = (Int32, Int32) -> Unit

///|

///|
/// The per-record half of `validate_records`.
fn validate_records_per_record(
  records : Array[Record],
  require_embedding~ : Bool = true,
) -> Unit raise ChromaError {
  if records.length() == 0 {
    raise ChromaError::InvalidArgument(
      "records must not be empty; Chroma rejects a batch with no records",
    )
  }
  let mut index = 0
  for record in records {
    if record.id == "" {
      raise ChromaError::InvalidArgument(
        "record at index " + index.to_string() + " has an empty id",
      )
    }
    if require_embedding && record.embedding is None {
      if record.document is None {
        raise ChromaError::InvalidArgument(
          "record " +
          record.id +
          " has neither a document nor an embedding; Chroma cannot store it",
        )
      }
      raise ChromaError::InvalidArgument(
        "record " +
        record.id +
        " has a document but no embedding; Chroma 1.x does not embed documents server-side, " +
        "so attach a vector with Record::with_embedding before storing it",
      )
    }
    index = index + 1
  }
}

///|
/// Rejects a query whose inputs would be refused server-side.
///
/// `label` names the field in the error message so a caller passing both
/// texts and vectors can tell which one was empty.
pub fn validate_query_inputs(
  label : String,
  inputs : Array[Json],
) -> Unit raise ChromaError {
  if inputs.length() == 0 {
    raise ChromaError::InvalidArgument(
      label + " must not be empty; a query needs at least one input",
    )
  }
}

///|
/// Builds a `where` clause for a metadata equality test.
///
/// The counterpart of chroma-go's `EqString(key, value)`: the clause filters
/// records whose metadata field `key` equals `value`, and can be combined
/// with `where_and` / `where_or`.
pub fn where_eq(key : String, value : MetadataValue) -> Json {
  let inner : Map[String, Json] = Map([])
  inner["$eq"] = value.to_json()
  let outer : Map[String, Json] = Map([])
  outer[key] = Json::object(inner)
  Json::object(outer)
}

///|
/// Builds a `where` clause for a metadata inequality test.
pub fn where_ne(key : String, value : MetadataValue) -> Json {
  let inner : Map[String, Json] = Map([])
  inner["$ne"] = value.to_json()
  let outer : Map[String, Json] = Map([])
  outer[key] = Json::object(inner)
  Json::object(outer)
}

///|
/// Builds a `where` clause matching records whose metadata field is greater
/// than `value`.
pub fn where_gt(key : String, value : MetadataValue) -> Json {
  comparison_clause(key, "$gt", value)
}

///|
/// Builds a `where` clause matching records whose metadata field is greater
/// than or equal to `value`.
pub fn where_gte(key : String, value : MetadataValue) -> Json {
  comparison_clause(key, "$gte", value)
}

///|
/// Builds a `where` clause matching records whose metadata field is less
/// than `value`.
pub fn where_lt(key : String, value : MetadataValue) -> Json {
  comparison_clause(key, "$lt", value)
}

///|
/// Builds a `where` clause matching records whose metadata field is less
/// than or equal to `value`.
pub fn where_lte(key : String, value : MetadataValue) -> Json {
  comparison_clause(key, "$lte", value)
}

///|
/// Builds a `where` clause matching records whose metadata field is one of
/// `values`.
pub fn where_in(key : String, values : Array[MetadataValue]) -> Json {
  let items : Array[Json] = []
  for value in values {
    items.push(value.to_json())
  }
  let inner : Map[String, Json] = Map([])
  inner["$in"] = Json::array(items)
  let outer : Map[String, Json] = Map([])
  outer[key] = Json::object(inner)
  Json::object(outer)
}

///|
fn comparison_clause(
  key : String,
  operator : String,
  value : MetadataValue,
) -> Json {
  let inner : Map[String, Json] = Map([])
  inner[operator] = value.to_json()
  let outer : Map[String, Json] = Map([])
  outer[key] = Json::object(inner)
  Json::object(outer)
}

///|
/// Combines clauses with `$and`, the counterpart of chroma-go's `And(...)`.
pub fn where_and(clauses : Array[Json]) -> Json {
  Json::object({ "$and": Json::array(clauses) })
}

///|
/// Combines clauses with `$or`, the counterpart of chroma-go's `Or(...)`.
pub fn where_or(clauses : Array[Json]) -> Json {
  Json::object({ "$or": Json::array(clauses) })
}

///|
/// Builds a `where_document` clause matching documents that contain `text`.
///
/// The counterpart of chroma-go's `Contains(text)`.
pub fn where_document_contains(text : String) -> Json {
  Json::object({ "$contains": Json::string(text) })
}

///|
/// Builds a `where_document` clause matching documents that do not contain
/// `text`.
pub fn where_document_not_contains(text : String) -> Json {
  Json::object({ "$not_contains": Json::string(text) })
}

///|
/// Builds a `where_document` clause matching documents by regular
/// expression.
pub fn where_document_regex(pattern : String) -> Json {
  Json::object({ "$regex": Json::string(pattern) })
}

///|
/// Combines `where_document` clauses with `$and`.
pub fn where_document_and(clauses : Array[Json]) -> Json {
  Json::object({ "$and": Json::array(clauses) })
}

///|
/// Combines `where_document` clauses with `$or`.
pub fn where_document_or(clauses : Array[Json]) -> Json {
  Json::object({ "$or": Json::array(clauses) })
}