// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// 移植自 milvus-io/milvus client/column/columns.go 的
// `Column.FieldData()` 一族与 client/milvusclient/write_options.go 的
// `processInsertColumns`（Apache-2.0）。
//
// 回读侧 `@column.from_field_data` 走的是 `FieldData` → `Column`；
// 这里走反方向 `Column` → `FieldData`。两边共用 `@column.ColumnValue`，
// 所以列类型只在 `column/` 定义一次 —— 写侧多出来的只有「行数对齐」
// 与「向量序列化」。

///|
/// 写入用的一列。名字加数据，等价于上游 `column.Column` 里写入侧用到的那部分。
///
/// 复用 `@column.ColumnValue` 而不是另造一套枚举，是为了让回读的
/// `Column` 能被原样改一改名字写回去（改列名重插入这种用法）。
pub(all) struct WriteColumn {
  name : String
  value : @column.ColumnValue
}

///|
pub fn WriteColumn::new(
  name : String,
  value : @column.ColumnValue,
) -> WriteColumn {
  { name, value, }
}

///|
/// 这一列的行数。
pub fn WriteColumn::len(self : WriteColumn) -> Int32 {
  match self.value {
    Bool(v) => v.length()
    Int8(v) | Int16(v) | Int32(v) => v.length()
    Int64(v) => v.length()
    Float(v) => v.length()
    Double(v) => v.length()
    String(v) | VarChar(v) | Text(v) => v.length()
    Timestamptz(v) => v.length()
    Json(v) => v.length()
    Geometry(v) => v.length()
    Array(v) => v.length()
    FloatVector(_, rows) | Float16Vector(_, rows) | BFloat16Vector(_, rows) =>
      rows.length()
    BinaryVector(_, rows) | Int8Vector(_, rows) => rows.length()
    SparseFloatVector(rows) => rows.length()
  }
}

///|
/// 这一列的 `@schema.DataType`。
pub fn WriteColumn::data_type(self : WriteColumn) -> @entity.DataType {
  match self.value {
    Bool(_) => @entity.DataType::Bool
    Int8(_) => @entity.DataType::Int8
    Int16(_) => @entity.DataType::Int16
    Int32(_) => @entity.DataType::Int32
    Int64(_) => @entity.DataType::Int64
    Float(_) => @entity.DataType::Float
    Double(_) => @entity.DataType::Double
    String(_) => @entity.DataType::String
    VarChar(_) => @entity.DataType::VarChar
    Text(_) => @entity.DataType::Text
    Timestamptz(_) => @entity.DataType::Timestamptz
    Json(_) => @entity.DataType::Json
    Geometry(_) => @entity.DataType::Geometry
    Array(_) => @entity.DataType::Array
    FloatVector(_, _) => @entity.DataType::FloatVector
    Float16Vector(_, _) => @entity.DataType::Float16Vector
    BFloat16Vector(_, _) => @entity.DataType::BFloat16Vector
    BinaryVector(_, _) => @entity.DataType::BinaryVector
    Int8Vector(_, _) => @entity.DataType::Int8Vector
    SparseFloatVector(_) => @entity.DataType::SparseFloatVector
  }
}

///|
/// 一列 → 一条 `FieldData`。
///
/// 动态字段（JSON）要带上 `is_dynamic`，否则服务端会把它当成 schema 里
/// 不存在的列拒绝。这里靠名字判断：`@column` 侧回读动态字段时也是用
/// `FieldData.is_dynamic` 标记的，写侧没有别的信号可用。
///
/// `nullable` 为真时写一条全 1 的 `valid_data`：调用方给的是实体数组，
/// 目前没有「某行是 null」的表示，写成全 1 与「全部有值」等价，
/// 但显式声明可空位图后服务端不会把空列当坏数据。
pub fn WriteColumn::to_field_data(
  self : WriteColumn,
  is_dynamic~ : Bool = false,
) -> @schema.FieldData raise ClientError {
  let rows = self.len()
  let field = match self.value {
    Bool(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(
          @schema.ScalarField_Data::BoolData(@schema.BoolArray::BoolArray(v)),
        ),
      )
    Int8(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(
          @schema.ScalarField_Data::IntData(
            @schema.IntArray::IntArray(v.map(narrow_int)),
          ),
        ),
      )
    Int16(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(
          @schema.ScalarField_Data::IntData(
            @schema.IntArray::IntArray(v.map(narrow_int)),
          ),
        ),
      )
    Int32(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(@schema.ScalarField_Data::IntData(@schema.IntArray::IntArray(v))),
      )
    Int64(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(
          @schema.ScalarField_Data::LongData(@schema.LongArray::LongArray(v)),
        ),
      )
    Float(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(
          @schema.ScalarField_Data::FloatData(@schema.FloatArray::FloatArray(v)),
        ),
      )
    Double(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(
          @schema.ScalarField_Data::DoubleData(
            @schema.DoubleArray::DoubleArray(v),
          ),
        ),
      )
    String(v) | VarChar(v) | Text(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(
          @schema.ScalarField_Data::StringData(
            @schema.StringArray::StringArray(v),
          ),
        ),
      )
    // JSON 在 wire 上是 `bytes`，每行一段 UTF-8 文本。
    Json(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(
          @schema.ScalarField_Data::JsonData(@schema.JSONArray::JSONArray(v)),
        ),
      )
    Timestamptz(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(
          @schema.ScalarField_Data::TimestamptzData(
            @schema.TimestamptzArray::TimestamptzArray(v),
          ),
        ),
      )
    Geometry(v) =>
      @schema.FieldData_Field::Scalars(
        scalar(
          @schema.ScalarField_Data::GeometryWktData(
            @schema.GeometryWktArray::GeometryWktArray(v),
          ),
        ),
      )
    Array(_) =>
      raise ClientError::Encode(
        "column \{self.name}: Array fields need an element type, which WriteColumn cannot infer",
      )
    // 浮点向量：一行 `Float` 数量就是宽度。
    FloatVector(dim, rows) =>
      vector_row_payload(
        self.name,
        dim,
        rows.map(row => row.length()),
        "float",
        () => {
          @schema.VectorField_Data::FloatVector(
            @schema.FloatArray::FloatArray(flatten(rows)),
          )
        },
      )
    Float16Vector(dim, rows) =>
      vector_row_payload(
        self.name,
        dim,
        rows.map(row => row.length()),
        "float16",
        () => {
          @schema.VectorField_Data::Float16Vector(
            bytes_of_rows(rows, @entity.float16_vector_bytes),
          )
        },
      )
    BFloat16Vector(dim, rows) =>
      vector_row_payload(
        self.name,
        dim,
        rows.map(row => row.length()),
        "bfloat16",
        () => {
          @schema.VectorField_Data::Bfloat16Vector(
            bytes_of_rows(rows, @entity.bfloat16_vector_bytes),
          )
        },
      )
    // 二进制向量：`dim` 是 **bit** 数，一行 `row.length()` 个字节管 8 倍。
    BinaryVector(dim, rows) =>
      vector_row_payload(
        self.name,
        dim,
        rows.map(row => row.length() * 8),
        "binary",
        () => @schema.VectorField_Data::BinaryVector(join_bytes(rows)),
      )
    // Int8 向量：`dim` 是 byte 数，与行 `Bytes` 长度同一套单位。
    Int8Vector(dim, rows) =>
      vector_row_payload(
        self.name,
        dim,
        rows.map(row => row.length()),
        "int8",
        () => @schema.VectorField_Data::Int8Vector(join_bytes(rows)),
      )
    SparseFloatVector(rows) =>
      @schema.FieldData_Field::Vectors(@schema.VectorField::{
        dim: 0L,
        valid_data: [],
        data: @schema.VectorField_Data::SparseFloatVector(@schema.SparseFloatArray::{
          contents: rows.map(row => {
            row.to_bytes() catch {
              err => raise ClientError::Schema(entity_schema_message(err))
            }
          }),
          dim: max_sparse_dim(rows).to_int64(),
        }),
      })
  }
  @schema.FieldData::{
    type_: data_type_to_proto(self.data_type()),
    field_name: self.name,
    field_id: 0L,
    is_dynamic,
    valid_data: [],
    field,
  }
  |> ignore_row_count(rows)
}

///|
/// 行数为 0 的列没有可写的内容，但也不该静默通过：
/// 服务端拿到空 `FieldData` 会把整批算成 0 行，调用方多半是搞错了。
fn ignore_row_count(
  field : @schema.FieldData,
  rows : Int32,
) -> @schema.FieldData raise ClientError {
  if rows == 0 {
    raise ClientError::Encode(
      "column \{field.field_name} has no rows; every column in a write must carry at least one",
    )
  }
  field
}

///|
/// 标量 payload：只填 data，`valid_data` 留空表示「全部有值」。
fn scalar(data : @schema.ScalarField_Data) -> @schema.ScalarField {
  @schema.ScalarField::{ valid_data: [], data, }
}

///|
/// 向量 payload 的公共外壳，顺带校一下 dim、行数与各行宽度。
///
/// 宽度校验放在 `build` 之前：报错要来自这里，而不是等 `flatten` 把
/// 不等宽的行拼成一段对不上行的字节流。
///
/// `widths` 是每行的「逻辑宽度」而不是 `Array` 的长度：二进制向量一行
/// 一个 byte 管 8 个 bit，`dim` 说的是 bit 数，两边得换算到同一套单位上。
fn vector_row_payload(
  name : String,
  dim : Int32,
  widths : Array[Int32],
  kind : String,
  build : () -> @schema.VectorField_Data raise ClientError,
) -> @schema.FieldData_Field raise ClientError {
  if dim <= 0 {
    raise ClientError::Encode(
      "column \{name}: vector dim must be positive, got \{dim}",
    )
  }
  if widths.length() == 0 {
    raise ClientError::Encode("column \{name}: vector column has no rows")
  }
  check_row_width(name, kind, dim, widths)
  @schema.FieldData_Field::Vectors(@schema.VectorField::{
    dim: dim.to_int64(),
    valid_data: [],
    data: build(),
  })
}

///|
/// 校一下「每行宽度一致」。
///
/// Milvus 是按 `dim` 这个**列级宽度**切值的：客户端的行边界它看不见，只有
/// 「值的总数」和「`num_rows`」。所以宽度对不上时，服务端会拿总数去除以 `dim`
/// 自己算出一个行数，再跟请求里的 `num_rows` 对不上就拒。真 v3.0.2 上实测：
///
/// ```
/// 单行 8 值 dim=4   -> the num_rows (2) of field (vector) is not equal
///                      to passed num_rows (1)
/// 两行各 6 值 dim=4 -> the num_rows (3) of field (vector) is not equal
///                      to passed num_rows (2)
/// 单行 3 值 dim=4   -> the length(3) of float data should divide the dim(4)
/// ```
///
/// 三条都不是「这一行有几列不对」，而是从 `dim` 反推出来的行数。其中最坏的
/// 一条是「值的总数能被 dim 整除、但行数与声明对不上」：这个形态下服务端会
/// **按 dim 重新切开**（8 个值被当成 2 行），只是最后被 `num_rows` 校验挡在
/// 门外。挡得住不代表这条路径是对的 —— 客户端比服务端多知道一件事：
/// 这一行本该有几列。在本地报出来，错误文本才指得准。
///
/// 判据是「各行宽度全等」，**不是「等于 dim」**。`FloatVector(dim, rows)`
/// 这个建模同时服务写入与查询向量，而查询向量的 `dim` 字段是冗余的
/// （每个元素就是一次查询，见 `client/search.mbt` 的契约）：建表 dim = 8、
/// 调用方按 4 写是合法用法，拿 dim 去卡会把本来正确的请求拦下来。
///
/// 单行时没有「各行」可言，`dim` 就是不可歧义的行宽，顺带对齐。
fn check_row_width(
  name : String,
  kind : String,
  dim : Int32,
  widths : Array[Int32],
) -> Unit raise ClientError {
  let first = widths[0]
  let mut uniform = true
  for width in widths {
    if width != first {
      uniform = false
    }
  }
  if !uniform {
    let detail : Array[String] = []
    for width in widths {
      detail.push(width.to_string())
    }
    raise ClientError::Encode(
      "column " +
      name +
      ": " +
      kind +
      " vector rows have different widths (first row has " +
      first.to_string() +
      " values, all rows: " +
      detail.join(",") +
      "); the server re-splits by dim " +
      dim.to_string() +
      " and the values silently go to the wrong rows",
    )
  }
  if widths.length() == 1 && first != dim {
    raise ClientError::Encode(
      "column " +
      name +
      ": " +
      kind +
      " vector row has " +
      first.to_string() +
      " values but dim is " +
      dim.to_string() +
      "; the server splits the payload by dim, not by row",
    )
  }
}

///|
/// 逐行编码后拼成一整块字节。fp16 / bf16 在 wire 上是「所有行首尾相接」，
/// 不是每行一个 bytes，所以这里必须拼平。
fn bytes_of_rows(
  rows : Array[Array[Float]],
  encode : (Array[Float]) -> Bytes,
) -> Bytes {
  join_bytes(rows.map(encode))
}

///|
fn join_bytes(parts : Array[Bytes]) -> Bytes {
  let buf = Buffer()
  for part in parts {
    buf.write_bytes(part)
  }
  buf.to_bytes()
}

///|
fn flatten(rows : Array[Array[Float]]) -> Array[Float] {
  let values : Array[Float] = []
  for row in rows {
    values.append(row)
  }
  values
}

///|
/// 稀疏向量那一列的 `dim` 是「本批次里最大的下标 + 1」，
/// 与上游 `SparseFloatArray.dim` 的语义一致（注释写明是最大维度）。
fn max_sparse_dim(rows : Array[@entity.SparseFloatVector]) -> UInt32 {
  let mut max : UInt32 = 0
  for row in rows {
    for index in row.indices {
      if index >= max {
        max = index + 1U
      }
    }
  }
  max
}

///|
/// `Int` → `Int32` 的窄化。越界的值直接报错，
/// 不静默按低位截断 —— 写进去一个错的值比写失败更难查。
fn narrow_int(value : Int32) -> Int32 raise ClientError {
  if value < -2147483648 || value > 2147483647 {
    raise ClientError::Encode(
      "value \{value} does not fit in a 32-bit integer field",
    )
  }
  value
}