// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors

import {
  Data as ArrowData,
  Table as ArrowTable,
  Binary,
  BufferType,
  DataType,
  DateUnit,
  Date_,
  Decimal,
  Dictionary,
  Duration,
  Field,
  FixedSizeBinary,
  FixedSizeList,
  Float,
  Float32,
  Int,
  LargeBinary,
  List,
  Null,
  Precision,
  RecordBatch,
  RecordBatchFileReader,
  RecordBatchFileWriter,
  RecordBatchStreamWriter,
  Schema,
  Struct,
  Timestamp,
  Type,
  Utf8,
  Vector,
  makeVector as arrowMakeVector,
  util as arrowUtil,
  vectorFromArray as badVectorFromArray,
  makeBuilder,
  makeData,
} from "apache-arrow";
import { Buffers } from "apache-arrow/data";
import { typedArrayToArrowType } from "./arrow_type";
import {
  blobToRead,
  coerceBlobValue,
  isBlobField,
  withBlobBytes,
} from "./blob";
import { type EmbeddingFunction } from "./embedding/embedding_function";
import {
  EmbeddingFunctionConfig,
  getRegistry,
  parseEmbeddingMetadata,
} from "./embedding/registry";
import {
  sanitizeField,
  sanitizeSchema,
  sanitizeTable,
  sanitizeType,
} from "./sanitize";
import { inferSchema } from "./schema";

export * from "apache-arrow";
export type SchemaLike =
  | Schema
  | {
      fields: FieldLike[];
      metadata: Map<string, string>;
      get names(): unknown[];
    };
export type FieldLike =
  | Field
  | {
      type: string | DataTypeLike;
      name: string;
      nullable: boolean;
      metadata?: Map<string, string>;
    };
/**
 * A `DataType` from any copy or version of apache-arrow.
 *
 * Arrow 21 brands its classes with `unique symbol` properties, so a type
 * object from a second copy of the library no longer satisfies the `DataType`
 * type of this one even though it is structurally identical. Inputs that only
 * need to be sanitized accept this looser shape instead.
 */
export type DataTypeLike = DataType | { readonly typeId: number };

/**
 * Create an Arrow field backed by LanceDB's JSON extension type.
 *
 * @param name - The field name.
 * @param nullable - Whether the field accepts null values.
 * @example
 * ```ts
 * import { connect, makeJsonField } from "@lancedb/lancedb";
 * import { Schema } from "apache-arrow";
 *
 * const schema = new Schema([makeJsonField("metadata")]);
 * const db = await connect("/path/to/database");
 * await db.createTable("items", [{ metadata: '{"source":"api"}' }], { schema });
 * ```
 */
export function makeJsonField(name: string, nullable = true): Field {
  return new Field(
    name,
    new Utf8(),
    nullable,
    new Map([["ARROW:extension:name", "arrow.json"]]),
  );
}

export type DataLike =
  | import("apache-arrow").Data
  | {
      // biome-ignore lint/suspicious/noExplicitAny: <explanation>
      type: any;
      length: number;
      offset: number;
      stride: number;
      nullable: boolean;
      children: DataLike[];
      dictionary?: { data: readonly DataLike[] };
      get nullCount(): number;
      // biome-ignore lint/suspicious/noExplicitAny: <explanation>
      values: Buffers<any>[BufferType.DATA];
      // biome-ignore lint/suspicious/noExplicitAny: <explanation>
      typeIds: Buffers<any>[BufferType.TYPE];
      // biome-ignore lint/suspicious/noExplicitAny: <explanation>
      nullBitmap: Buffers<any>[BufferType.VALIDITY];
      // biome-ignore lint/suspicious/noExplicitAny: <explanation>
      valueOffsets: Buffers<any>[BufferType.OFFSET];
    };

export type RecordBatchLike =
  | RecordBatch
  | {
      schema: SchemaLike;
      data: DataLike;
    };

export type TableLike =
  | ArrowTable
  | { schema: SchemaLike; batches: RecordBatchLike[] };

export type IntoVector =
  | Float32Array
  | Float64Array
  | Uint8Array
  | number[]
  | Promise<Float32Array | Float64Array | Uint8Array | number[]>;

export type MultiVector = IntoVector[];

export function isMultiVector(value: unknown): value is MultiVector {
  return Array.isArray(value) && isIntoVector(value[0]);
}

// Float16Array is not in TypeScript's standard lib yet; access dynamically
type Float16ArrayCtor = new (
  ...args: unknown[]
) => { buffer: ArrayBuffer; byteOffset: number; byteLength: number };
const float16ArrayCtor = (globalThis as unknown as Record<string, unknown>)
  .Float16Array as Float16ArrayCtor | undefined;

export function isIntoVector(value: unknown): value is IntoVector {
  return (
    value instanceof Float32Array ||
    value instanceof Float64Array ||
    value instanceof Uint8Array ||
    (float16ArrayCtor !== undefined && value instanceof float16ArrayCtor) ||
    (Array.isArray(value) && !Array.isArray(value[0]))
  );
}

/**
 * Extract the underlying byte buffer and data type from a typed array
 * for passing to the Rust NAPI layer without precision loss.
 */
export function extractVectorBuffer(
  vector: Float32Array | Float64Array | Uint8Array,
): { data: Uint8Array; dtype: string } | null {
  if (float16ArrayCtor !== undefined && vector instanceof float16ArrayCtor) {
    return {
      data: new Uint8Array(vector.buffer, vector.byteOffset, vector.byteLength),
      dtype: "float16",
    };
  }
  if (vector instanceof Float64Array) {
    return {
      data: new Uint8Array(vector.buffer, vector.byteOffset, vector.byteLength),
      dtype: "float64",
    };
  }
  if (vector instanceof Uint8Array && !(vector instanceof Float32Array)) {
    return { data: vector, dtype: "uint8" };
  }
  return null;
}

export function isArrowTable(value: object): value is TableLike {
  if (value instanceof ArrowTable) return true;
  return "schema" in value && "batches" in value;
}

export function isNull(value: unknown): value is Null {
  return value instanceof Null || DataType.isNull(value);
}
export function isInt(value: unknown): value is Int {
  return value instanceof Int || DataType.isInt(value);
}
export function isFloat(value: unknown): value is Float {
  return value instanceof Float || DataType.isFloat(value);
}
export function isBinary(value: unknown): value is Binary {
  return value instanceof Binary || DataType.isBinary(value);
}
export function isLargeBinary(value: unknown): value is LargeBinary {
  return value instanceof LargeBinary || DataType.isLargeBinary(value);
}
export function isUtf8(value: unknown): value is Utf8 {
  return value instanceof Utf8 || DataType.isUtf8(value);
}
export function isLargeUtf8(value: unknown): value is Utf8 {
  return value instanceof Utf8 || DataType.isLargeUtf8(value);
}
export function isBool(value: unknown): value is Utf8 {
  return value instanceof Utf8 || DataType.isBool(value);
}
export function isDecimal(value: unknown): value is Utf8 {
  return value instanceof Utf8 || DataType.isDecimal(value);
}
export function isDate(value: unknown): value is Utf8 {
  return value instanceof Utf8 || DataType.isDate(value);
}
export function isTime(value: unknown): value is Utf8 {
  return value instanceof Utf8 || DataType.isTime(value);
}
export function isTimestamp(value: unknown): value is Utf8 {
  return value instanceof Utf8 || DataType.isTimestamp(value);
}
export function isInterval(value: unknown): value is Utf8 {
  return value instanceof Utf8 || DataType.isInterval(value);
}
export function isDuration(value: unknown): value is Utf8 {
  return value instanceof Utf8 || DataType.isDuration(value);
}
export function isList(value: unknown): value is List {
  return value instanceof List || DataType.isList(value);
}
export function isStruct(value: unknown): value is Struct {
  return value instanceof Struct || DataType.isStruct(value);
}
export function isUnion(value: unknown): value is Struct {
  return value instanceof Struct || DataType.isUnion(value);
}
export function isFixedSizeBinary(value: unknown): value is FixedSizeBinary {
  return value instanceof FixedSizeBinary || DataType.isFixedSizeBinary(value);
}

export function isFixedSizeList(value: unknown): value is FixedSizeList {
  return value instanceof FixedSizeList || DataType.isFixedSizeList(value);
}

/** Data type accepted by NodeJS SDK */
export type Data = Record<string, unknown>[] | TableLike;

/*
 * Options to control how a column should be converted to a vector array
 */
export class VectorColumnOptions {
  /** Vector column type. */
  type: Float = new Float32();

  constructor(values?: Partial<VectorColumnOptions>) {
    Object.assign(this, values);
  }
}

// biome-ignore lint/suspicious/noExplicitAny: skip
function vectorFromArray(data: any, type?: DataType) {
  // Workaround for: https://github.com/apache/arrow/issues/45862
  // If FSL type with float
  if (DataType.isFixedSizeList(type) && DataType.isFloat(type.valueType)) {
    const extendedData = [...data, new Array(type.listSize).fill(0.0)];
    const array = badVectorFromArray(extendedData, type);
    return array.slice(0, data.length);
  } else if (type === undefined) {
    return badVectorFromArray(data);
  } else {
    return badVectorFromArray(data, type);
  }
}

/** Options to control the makeArrowTable call. */
export class MakeArrowTableOptions {
  /*
   * Schema of the data.
   *
   * If this is not provided then the data type will be inferred from the
   * JS type.  Integer numbers will become int64, floating point numbers
   * will become float64 and arrays will become variable sized lists with
   * the data type inferred from the first element in the array.
   *
   * The schema must be specified if there are no records (e.g. to make
   * an empty table)
   */
  schema?: SchemaLike;

  /*
   * Mapping from vector column name to expected type
   *
   * Lance expects vector columns to be fixed size list arrays (i.e. tensors)
   * However, `makeArrowTable` will not infer this by default (it creates
   * variable size list arrays).  This field can be used to indicate that a column
   * should be treated as a vector column and converted to a fixed size list.
   *
   * The keys should be the names of the vector columns.  The value specifies the
   * expected data type of the vector columns.
   *
   * If `schema` is provided then this field is ignored.
   *
   * By default, the column named "vector" will be assumed to be a float32
   * vector column.
   */
  vectorColumns: Record<string, VectorColumnOptions> = {
    vector: new VectorColumnOptions(),
  };
  embeddings?: EmbeddingFunction<unknown>;
  embeddingFunction?: EmbeddingFunctionConfig;

  /**
   * If true then string columns will be encoded with dictionary encoding
   *
   * Set this to true if your string columns tend to repeat the same values
   * often.  For more precise control use the `schema` property to specify the
   * data type for individual columns.
   *
   * If `schema` is provided then this property is ignored.
   */
  dictionaryEncodeStrings: boolean = false;

  constructor(values?: Partial<MakeArrowTableOptions>) {
    Object.assign(this, values);
  }
}

/**
 * An enhanced version of the apache-arrow makeTable function from Apache Arrow
 * that supports nested fields and embeddings columns.
 *
 * (typically you do not need to call this function.  It will be called automatically
 * when creating a table or adding data to it)
 *
 * This function converts an array of Record<String, any> (row-major JS objects)
 * to an Arrow Table (a columnar structure)
 *
 * If a schema is provided then it will be used to determine the resulting array
 * types.  Fields will also be reordered to fit the order defined by the schema.
 *
 * If a schema is not provided then the types will be inferred and the field order
 * will be controlled by the order of properties in the first record.  If a type
 * is inferred it will always be nullable.
 *
 * If not all fields are found in the data, then a subset of the schema will be
 * returned.
 *
 * If the input is empty then a schema must be provided to create an empty table.
 *
 * When a schema is not specified then data types will be inferred.  The inference
 * rules are as follows:
 *
 *  - boolean => Bool
 *  - number => Float64
 *  - bigint => Int64
 *  - String => Utf8
 *  - Buffer => Binary
 *  - Record<String, any> => Struct
 *  - Array<any> => List
 * @example
 * ```ts
 * import { fromTableToBuffer, makeArrowTable } from "../arrow";
 * import { Field, FixedSizeList, Float16, Float32, Int32, Schema } from "apache-arrow";
 *
 * const schema = new Schema([
 *   new Field("a", new Int32()),
 *   new Field("b", new Float32()),
 *   new Field("c", new FixedSizeList(3, new Field("item", new Float16()))),
 *  ]);
 *  const table = makeArrowTable([
 *    { a: 1, b: 2, c: [1, 2, 3] },
 *    { a: 4, b: 5, c: [4, 5, 6] },
 *    { a: 7, b: 8, c: [7, 8, 9] },
 *  ], { schema });
 * ```
 *
 * By default it assumes that the column named `vector` is a vector column
 * and it will be converted into a fixed size list array of type float32.
 * The `vectorColumns` option can be used to support other vector column
 * names and data types.
 *
 * ```ts
 * const schema = new Schema([
 *   new Field("a", new Float64()),
 *   new Field("b", new Float64()),
 *   new Field(
 *     "vector",
 *     new FixedSizeList(3, new Field("item", new Float32()))
 *   ),
 * ]);
 * const table = makeArrowTable([
 *   { a: 1, b: 2, vector: [1, 2, 3] },
 *   { a: 4, b: 5, vector: [4, 5, 6] },
 *   { a: 7, b: 8, vector: [7, 8, 9] },
 * ]);
 * assert.deepEqual(table.schema, schema);
 * ```
 *
 * You can specify the vector column types and names using the options as well
 *
 * ```ts
 * const schema = new Schema([
 *   new Field('a', new Float64()),
 *   new Field('b', new Float64()),
 *   new Field('vec1', new FixedSizeList(3, new Field('item', new Float16()))),
 *   new Field('vec2', new FixedSizeList(3, new Field('item', new Float16())))
 * ]);
 * const table = makeArrowTable([
 *   { a: 1, b: 2, vec1: [1, 2, 3], vec2: [2, 4, 6] },
 *   { a: 4, b: 5, vec1: [4, 5, 6], vec2: [8, 10, 12] },
 *   { a: 7, b: 8, vec1: [7, 8, 9], vec2: [14, 16, 18] }
 * ], {
 *   vectorColumns: {
 *     vec1: { type: new Float16() },
 *     vec2: { type: new Float16() }
 *   }
 * }
 * assert.deepEqual(table.schema, schema)
 * ```
 */
export function makeArrowTable(
  data: Array<Record<string, unknown>>,
  options?: Partial<MakeArrowTableOptions>,
  metadata?: Map<string, string>,
): ArrowTable {
  const opt = new MakeArrowTableOptions(options !== undefined ? options : {});
  let schema: Schema | undefined = undefined;
  if (opt.schema !== undefined && opt.schema !== null) {
    schema = sanitizeSchema(opt.schema);
    schema = validateSchemaEmbeddings(
      schema as Schema,
      data,
      options?.embeddingFunction,
    );
  }

  if (schema !== undefined) {
    // Validate blob values up front and give every one the full
    // `{ data, uri }` shape, so inference sees the same struct whether a row
    // passed bytes, a URI, or a partial struct.
    data = mapBlobInputs(data, schema, (value, field, row) => {
      if (value === undefined) {
        return value;
      }
      try {
        return coerceBlobValue(value);
      } catch (e) {
        throw new Error(
          `Invalid value for blob field ${field} at row ${row}: ${(e as Error).message}`,
        );
      }
    });
  }

  let schemaMetadata = schema?.metadata || new Map<string, string>();
  if (metadata !== undefined) {
    schemaMetadata = new Map([...schemaMetadata, ...metadata]);
  }

  if (
    data.length === 0 &&
    (options?.schema === undefined || options?.schema === null)
  ) {
    throw new Error("At least one record or a schema needs to be provided");
  } else if (data.length === 0) {
    if (schema === undefined) {
      throw new Error("A schema must be provided if data is empty");
    } else {
      schema = new Schema(schema.fields, schemaMetadata);
      validateBlobSchema(schema);
      return new ArrowTable(schema);
    }
  }

  let inferredSchema = inferSchema(data, schema, opt);
  inferredSchema = new Schema(inferredSchema.fields, schemaMetadata);
  validateBlobSchema(inferredSchema);

  const finalColumns: Record<string, Vector> = {};
  for (const field of inferredSchema.fields) {
    finalColumns[field.name] = transposeData(data, field);
  }

  return new ArrowTable(inferredSchema, finalColumns);
}

function validateBlobSchema(schema: Schema): void {
  for (const field of schema.fields) {
    validateBlobField(field);
  }
}

function validateBlobField(field: Field): void {
  if (
    isFixedSizeList(field.type) &&
    containsBlobField(field.type.children[0])
  ) {
    throw new Error(
      "Blob fields inside FixedSizeList are not supported. Use List instead.",
    );
  }
  for (const child of field.type.children ?? []) {
    validateBlobField(child);
  }
}

function containsBlobField(field: Field): boolean {
  if (isBlobField(field)) {
    return true;
  }
  return (field.type.children ?? []).some((child: Field) =>
    containsBlobField(child),
  );
}

type BlobInputVisitor = (value: unknown, field: string, row: number) => unknown;

/**
 * Calls `visit` on every value that lands in a blob column of `schema`,
 * including blobs nested in structs and lists, and replaces it with the
 * result. Records and arrays are copied only where a value changed, so the
 * caller's data is never mutated.
 */
function mapBlobInputs(
  data: Array<Record<string, unknown>>,
  schema: Schema,
  visit: BlobInputVisitor,
): Array<Record<string, unknown>> {
  if (!schema.fields.some(containsBlobField)) {
    return data;
  }
  return data.map((record, row) =>
    isObject(record)
      ? mapBlobFields(record, schema.fields, "", row, visit)
      : record,
  );
}

function mapBlobFields(
  record: Record<string, unknown>,
  fields: Field[],
  prefix: string,
  row: number,
  visit: BlobInputVisitor,
): Record<string, unknown> {
  let out: Record<string, unknown> | undefined;
  for (const field of fields) {
    if (!containsBlobField(field) || !Object.hasOwn(record, field.name)) {
      continue;
    }
    const value = record[field.name];
    const mapped = mapBlobField(
      value,
      field,
      `${prefix}${field.name}`,
      row,
      visit,
    );
    if (mapped !== value) {
      out ??= { ...record };
      out[field.name] = mapped;
    }
  }
  return out ?? record;
}

function mapBlobField(
  value: unknown,
  field: Field,
  label: string,
  row: number,
  visit: BlobInputVisitor,
): unknown {
  if (isBlobField(field)) {
    return visit(value, label, row);
  }
  if (field.type instanceof Struct && isObject(value)) {
    return mapBlobFields(value, field.type.children, `${label}.`, row, visit);
  }
  if (isList(field.type) && Array.isArray(value)) {
    const child = field.type.children[0];
    let out: unknown[] | undefined;
    for (const [index, element] of value.entries()) {
      const mapped = mapBlobField(
        element,
        child,
        `${label}[${index}]`,
        row,
        visit,
      );
      if (mapped !== element) {
        out ??= [...value];
        out[index] = mapped;
      }
    }
    return out ?? value;
  }
  return value;
}

/**
 * Reads `Blob` / `File` values in the blob columns of `schema` into bytes so
 * the synchronous conversion in {@link makeArrowTable} can accept them.
 * Returns `data` itself when there is nothing to read.
 */
export async function resolveBlobInputs(
  data: Array<Record<string, unknown>>,
  schema?: SchemaLike,
): Promise<Array<Record<string, unknown>>> {
  if (schema === undefined || schema === null) {
    return data;
  }
  const sanitized = sanitizeSchema(schema);
  // Keyed by the Blob itself, so rows sharing one Blob (bare or as
  // `{ data }`) read it once and share the bytes.
  const reads = new Map<Blob, Promise<Uint8Array>>();
  mapBlobInputs(data, sanitized, (value) => {
    const source = blobToRead(value);
    if (source !== undefined && !reads.has(source)) {
      reads.set(
        source,
        source.arrayBuffer().then((buffer) => new Uint8Array(buffer)),
      );
    }
    return value;
  });
  if (reads.size === 0) {
    return data;
  }
  const bytes = new Map(
    await Promise.all(
      [...reads].map(async ([source, read]) => [source, await read] as const),
    ),
  );
  return mapBlobInputs(data, sanitized, (value) => {
    const source = blobToRead(value);
    return source === undefined
      ? value
      : withBlobBytes(value, bytes.get(source) as Uint8Array);
  });
}

function isObject(value: unknown): value is Record<string, unknown> {
  return (
    typeof value === "object" &&
    value !== null &&
    !Array.isArray(value) &&
    !(value instanceof RegExp) &&
    !(value instanceof Date) &&
    !(value instanceof Set) &&
    !(value instanceof Map) &&
    !(value instanceof Buffer) &&
    !ArrayBuffer.isView(value)
  );
}

function valueAtPath(datum: Record<string, unknown>, path: string[]): unknown {
  let current: unknown = datum;
  for (const key of path) {
    if (current == null) {
      return null;
    }
    if (isObject(current) && (Object.hasOwn(current, key) || key in current)) {
      current = current[key];
    } else {
      return undefined;
    }
  }
  return current;
}

function transposeData(
  data: Record<string, unknown>[],
  field: Field,
  path: string[] = [],
): Vector {
  const valuesPath = [...path, field.name];
  if (isBlobField(field) && field.type instanceof Struct) {
    const blobRows = data.map((datum) =>
      coerceBlobValue(valueAtPath(datum, valuesPath)),
    );
    const childVectors = field.type.children.map((child) => {
      const values = blobRows.map((row) =>
        row == null ? null : (row[child.name as "data" | "uri"] ?? null),
      );
      return makeVector(values, child.type, undefined, child.nullable);
    });
    const nullCount = blobRows.filter((row) => row === null).length;
    const structData = makeData({
      type: field.type,
      length: blobRows.length,
      nullCount,
      nullBitmap:
        nullCount > 0
          ? arrowUtil.packBools(blobRows.map((row) => row !== null))
          : undefined,
      children: childVectors.map((v) => v.data[0]),
    });
    return arrowMakeVector(structData);
  }
  if (isList(field.type) && containsBlobField(field.type.children[0])) {
    return transposeListData(data, field, valuesPath);
  }
  const values = data.map((datum) => valueAtPath(datum, valuesPath));
  if (field.type instanceof Struct) {
    const childFields = field.type.children;
    const childVectors = childFields.map((child) => {
      return transposeData(data, child, valuesPath);
    });
    const nullCount = values.filter((value) => value === null).length;
    const structData = makeData({
      type: field.type,
      length: values.length,
      nullCount,
      nullBitmap:
        nullCount > 0
          ? arrowUtil.packBools(values.map((value) => value !== null))
          : undefined,
      children: childVectors.map((v) => v.data[0]),
    });
    return arrowMakeVector(structData);
  } else {
    return makeVector(values, field.type, undefined, field.nullable);
  }
}

function transposeListData(
  data: Record<string, unknown>[],
  field: Field,
  valuesPath: string[],
): Vector {
  const listType = field.type as List;
  const childField = listType.children[0];
  const lists = data.map((datum) => valueAtPath(datum, valuesPath));
  const flattened: Record<string, unknown>[] = [];
  const validity: boolean[] = [];
  const offsets: number[] = [0];

  for (const list of lists) {
    if (list == null) {
      validity.push(false);
      offsets.push(flattened.length);
      continue;
    }
    if (!Array.isArray(list)) {
      throw new Error(`expected an array for list field '${field.name}'`);
    }
    validity.push(true);
    for (const element of list) {
      flattened.push({ [childField.name]: element });
    }
    offsets.push(flattened.length);
  }

  const childVector = transposeData(flattened, childField, []);
  const nullCount = validity.filter((valid) => !valid).length;
  return arrowMakeVector(
    makeData({
      type: listType,
      length: lists.length,
      nullCount,
      nullBitmap: nullCount > 0 ? arrowUtil.packBools(validity) : undefined,
      valueOffsets: Int32Array.from(offsets),
      child: childVector.data[0],
    }),
  );
}

/**
 * Create an empty Arrow table with the provided schema
 */
export function makeEmptyTable(
  schema: SchemaLike,
  metadata?: Map<string, string>,
): ArrowTable {
  return makeArrowTable([], { schema }, metadata);
}

/**
 * Helper function to convert Array<Array<any>> to a variable sized list array
 */
// @ts-expect-error (Vector<unknown> is not assignable to Vector<any>)
function makeListVector(lists: unknown[][]): Vector<unknown> {
  if (lists.length === 0 || lists[0].length === 0) {
    throw Error("Cannot infer list vector from empty array or empty list");
  }
  const sampleList = lists[0];
  // biome-ignore lint/suspicious/noExplicitAny: skip
  let inferredType: any;
  try {
    const sampleVector = makeVector(sampleList);
    inferredType = sampleVector.type;
  } catch (error: unknown) {
    // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
    throw Error(`Cannot infer list vector.  Cannot infer inner type: ${error}`);
  }

  const listBuilder = makeBuilder({
    type: new List(new Field("item", inferredType, true)),
  });
  for (const list of lists) {
    listBuilder.append(list);
  }
  return listBuilder.finish().toVector();
}

/** Helper function to convert an Array of JS values to an Arrow Vector */
function makeVector(
  values: unknown[],
  type?: DataType,
  stringAsDictionary?: boolean,
  nullable?: boolean,
  // biome-ignore lint/suspicious/noExplicitAny: skip
): Vector<any> {
  if (type !== undefined) {
    // Convert undefined values to null for nullable fields
    if (nullable) {
      values = values.map((v) => (v === undefined ? null : v));
    }

    // workaround for: https://github.com/apache/arrow-js/issues/68
    if (DataType.isBool(type)) {
      const hasNonNullValue = values.some((v) => v !== null && v !== undefined);
      if (!hasNonNullValue) {
        const nullBitmap = new Uint8Array(Math.ceil(values.length / 8));
        const data = makeData({
          type: type,
          length: values.length,
          nullCount: values.length,
          nullBitmap,
        });
        return arrowMakeVector(data);
      }
    }

    // No need for inference, let Arrow create it
    if (type instanceof Int) {
      if (DataType.isInt(type) && type.bitWidth === 64) {
        // wrap in BigInt to avoid bug: https://github.com/apache/arrow/issues/40051
        values = values.map((v) => {
          if (v === null) {
            return v;
          } else if (typeof v === "bigint") {
            return v;
          } else if (typeof v === "number") {
            return BigInt(v);
          } else {
            return v;
          }
        });
      } else {
        // Similarly, bigint isn't supported for 16 or 32-bit ints.
        values = values.map((v) => {
          if (typeof v == "bigint") {
            return Number(v);
          } else {
            return v;
          }
        });
      }
    }
    return vectorFromArray(values, type);
  }
  if (values.length === 0) {
    throw Error(
      "makeVector requires at least one value or the type must be specified",
    );
  }
  const sampleValue = values.find((val) => val !== null && val !== undefined);
  if (sampleValue === undefined) {
    throw Error(
      "makeVector cannot infer the type if all values are null or undefined",
    );
  }
  if (ArrayBuffer.isView(sampleValue) && !(sampleValue instanceof DataView)) {
    const info = typedArrayToArrowType(sampleValue);
    if (info !== undefined) {
      const fslType = new FixedSizeList(
        info.length,
        new Field("item", info.elementType, true),
      );
      return vectorFromArray(values, fslType);
    }
  }
  if (Array.isArray(sampleValue)) {
    // Default Arrow inference doesn't handle list types
    return makeListVector(values as unknown[][]);
  } else if (Buffer.isBuffer(sampleValue)) {
    // Default Arrow inference doesn't handle Buffer
    return vectorFromArray(values, new Binary());
  } else if (
    !(stringAsDictionary ?? false) &&
    (typeof sampleValue === "string" || sampleValue instanceof String)
  ) {
    // If the type is string then don't use Arrow's default inference unless dictionaries are requested
    // because it will always use dictionary encoding for strings
    return vectorFromArray(values, new Utf8());
  } else {
    // Convert a JS array of values to an arrow vector
    return vectorFromArray(values);
  }
}

/** Helper function to apply embeddings from metadata to an input table */
async function applyEmbeddingsFromMetadata(
  table: ArrowTable,
  schema: Schema,
): Promise<ArrowTable> {
  const registry = getRegistry();
  const functions = await registry.parseFunctions(schema.metadata);

  const columns = Object.fromEntries(
    table.schema.fields.map((field) => [
      field.name,
      table.getChild(field.name)!,
    ]),
  );

  for (const functionEntry of functions.values()) {
    const sourceColumn = columns[functionEntry.sourceColumn];
    const destColumn = functionEntry.vectorColumn;
    if (sourceColumn === undefined) {
      throw new Error(
        `Cannot apply embedding function because the source column '${functionEntry.sourceColumn}' was not present in the data`,
      );
    }

    // Check if destination column exists and handle accordingly
    if (columns[destColumn] !== undefined) {
      const existingColumn = columns[destColumn];
      // If the column exists but is all null, we can fill it with embeddings
      if (existingColumn.nullCount !== existingColumn.length) {
        // Column has non-null values, skip embedding application
        continue;
      }
    }

    if (table.batches.length > 1) {
      throw new Error(
        "Internal error: `makeArrowTable` unexpectedly created a table with more than one batch",
      );
    }
    const values = sourceColumn.toArray();

    const vectors =
      await functionEntry.function.computeSourceEmbeddings(values);
    if (vectors.length !== values.length) {
      throw new Error(
        "Embedding function did not return an embedding for each input element",
      );
    }
    let destType: DataType;
    const dtype = schema.fields.find((f) => f.name === destColumn)!.type;
    if (isFixedSizeList(dtype)) {
      destType = sanitizeType(dtype);
    } else {
      throw new Error(
        "Expected FixedSizeList as datatype for vector field, instead got: " +
          dtype,
      );
    }
    const vector = makeVector(vectors, destType);
    columns[destColumn] = vector;
  }

  // Add any missing columns from the schema as null vectors
  for (const field of schema.fields) {
    if (!(field.name in columns)) {
      const nullValues = new Array(table.numRows).fill(null);
      columns[field.name] = makeVector(
        nullValues,
        field.type,
        undefined,
        field.nullable,
      );
    }
  }

  const newTable = new ArrowTable(columns);
  return alignTable(newTable, schema);
}

/** Helper function to apply embeddings to an input table */
async function applyEmbeddings<T>(
  table: ArrowTable,
  embeddings?: EmbeddingFunctionConfig,
  schema?: SchemaLike,
): Promise<ArrowTable> {
  if (schema !== undefined && schema !== null) {
    schema = sanitizeSchema(schema);
  }
  if (schema?.metadata.has("embedding_functions")) {
    return applyEmbeddingsFromMetadata(table, schema! as Schema);
  } else if (embeddings == null || embeddings === undefined) {
    return table;
  }

  let schemaMetadata = schema?.metadata || new Map<string, string>();

  if (!(embeddings == null || embeddings === undefined)) {
    const registry = getRegistry();
    const embeddingMetadata = registry.getTableMetadata([embeddings]);
    schemaMetadata = new Map([...schemaMetadata, ...embeddingMetadata]);
  }

  // Convert from ArrowTable to Record<String, Vector>
  const colEntries = [...Array(table.numCols).keys()].map((_, idx) => {
    const name = table.schema.fields[idx].name;
    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
    const vec = table.getChildAt(idx)!;
    return [name, vec];
  });
  const newColumns = Object.fromEntries(colEntries);

  const sourceColumn = newColumns[embeddings.sourceColumn];
  const destColumn = embeddings.vectorColumn ?? "vector";
  const innerDestType =
    embeddings.function.embeddingDataType() ?? new Float32();
  if (sourceColumn === undefined) {
    throw new Error(
      `Cannot apply embedding function because the source column '${embeddings.sourceColumn}' was not present in the data`,
    );
  }

  if (table.numRows === 0) {
    if (Object.prototype.hasOwnProperty.call(newColumns, destColumn)) {
      // We have an empty table and it already has the embedding column so no work needs to be done
      // Note: we don't return an error like we did below because this is a common occurrence.  For example,
      // if we call convertToTable with 0 records and a schema that includes the embedding
      return table;
    }
    const dimensions = embeddings.function.ndims();
    if (dimensions !== undefined) {
      const destType = newVectorType(dimensions, innerDestType);
      newColumns[destColumn] = makeVector([], destType);
    } else if (schema != null) {
      const destField = schema.fields.find((f) => f.name === destColumn);
      if (destField != null) {
        newColumns[destColumn] = makeVector(
          [],
          destField.type,
          undefined,
          destField.nullable,
        );
      } else {
        throw new Error(
          `Attempt to apply embeddings to an empty table failed because schema was missing embedding column '${destColumn}'`,
        );
      }
    } else {
      throw new Error(
        "Attempt to apply embeddings to an empty table when the embeddings function does not specify `embeddingDimension`",
      );
    }
  } else {
    // Check if destination column exists and handle accordingly
    if (Object.prototype.hasOwnProperty.call(newColumns, destColumn)) {
      const existingColumn = newColumns[destColumn];
      // If the column exists but is all null, we can fill it with embeddings
      if (existingColumn.nullCount !== existingColumn.length) {
        // Column has non-null values, skip embedding application and return table as-is
        let newTable = new ArrowTable(newColumns);
        if (schema != null) {
          newTable = alignTable(newTable, schema as Schema);
        }
        return new ArrowTable(
          new Schema(newTable.schema.fields, schemaMetadata),
          newTable.batches,
        );
      }
    }

    if (table.batches.length > 1) {
      throw new Error(
        "Internal error: `makeArrowTable` unexpectedly created a table with more than one batch",
      );
    }
    const values = sourceColumn.toArray();
    const vectors = await embeddings.function.computeSourceEmbeddings(
      values as T[],
    );
    if (vectors.length !== values.length) {
      throw new Error(
        "Embedding function did not return an embedding for each input element",
      );
    }
    const destType = newVectorType(vectors[0].length, innerDestType);
    newColumns[destColumn] = makeVector(vectors, destType);
  }

  let newTable = new ArrowTable(newColumns);
  if (schema != null) {
    if (schema.fields.find((f) => f.name === destColumn) === undefined) {
      throw new Error(
        `When using embedding functions and specifying a schema the schema should include the embedding column but the column ${destColumn} was missing`,
      );
    }
    newTable = alignTable(newTable, schema as Schema);
  }

  newTable = new ArrowTable(
    new Schema(newTable.schema.fields, schemaMetadata),
    newTable.batches,
  );

  return newTable;
}

/**
 * Convert an Array of records into an Arrow Table, optionally applying an
 * embeddings function to it.
 *
 * This function calls `makeArrowTable` first to create the Arrow Table.
 * Any provided `makeTableOptions` (e.g. a schema) will be passed on to
 * that call.
 *
 * The embedding function will be passed a column of values (based on the
 * `sourceColumn` of the embedding function) and expects to receive back
 * number[][] which will be converted into a fixed size list column.  By
 * default this will be a fixed size list of Float32 but that can be
 * customized by the `embeddingDataType` property of the embedding function.
 *
 * If a schema is provided in `makeTableOptions` then it should include the
 * embedding columns.  If no schema is provided then embedding columns will
 * be placed at the end of the table, after all of the input columns.
 */
export async function convertToTable(
  data: Array<Record<string, unknown>>,
  embeddings?: EmbeddingFunctionConfig,
  makeTableOptions?: Partial<MakeArrowTableOptions>,
): Promise<ArrowTable> {
  let processedData = data;

  // If we have a schema with embedding metadata, we need to preprocess the data
  // to ensure all nested fields are present
  if (
    makeTableOptions?.schema &&
    makeTableOptions.schema.metadata?.has("embedding_functions")
  ) {
    processedData = ensureNestedFieldsExist(
      data,
      makeTableOptions.schema as Schema,
    );
  }

  processedData = await resolveBlobInputs(
    processedData,
    makeTableOptions?.schema,
  );

  const table = makeArrowTable(processedData, makeTableOptions);
  return await applyEmbeddings(table, embeddings, makeTableOptions?.schema);
}

/** Creates the Arrow Type for a Vector column with dimension `dim` */
export function newVectorType<T extends Float>(
  dim: number,
  innerType: unknown,
): FixedSizeList<T> {
  // in Lance we always default to have the elements nullable, so we need to set it to true
  // otherwise we often get schema mismatches because the stored data always has schema with nullable elements
  const children = new Field("item", <T>sanitizeType(innerType), true);
  return new FixedSizeList(dim, children);
}

/**
 * Serialize an Array of records into a buffer using the Arrow IPC File serialization
 *
 * This function will call `convertToTable` and pass on `embeddings` and `schema`
 *
 * `schema` is required if data is empty
 */
export async function fromRecordsToBuffer(
  data: Array<Record<string, unknown>>,
  embeddings?: EmbeddingFunctionConfig,
  schema?: Schema,
): Promise<Buffer> {
  if (schema !== undefined && schema !== null) {
    schema = sanitizeSchema(schema);
  }
  const table = await convertToTable(data, embeddings, { schema });
  const writer = RecordBatchFileWriter.writeAll(table);
  return Buffer.from(await writer.toUint8Array());
}

/**
 * Serialize an Array of records into a buffer using the Arrow IPC Stream serialization
 *
 * This function will call `convertToTable` and pass on `embeddings` and `schema`
 *
 * `schema` is required if data is empty
 */
export async function fromRecordsToStreamBuffer(
  data: Array<Record<string, unknown>>,
  embeddings?: EmbeddingFunctionConfig,
  schema?: Schema,
): Promise<Buffer> {
  if (schema !== undefined && schema !== null) {
    schema = sanitizeSchema(schema);
  }
  const table = await convertToTable(data, embeddings, { schema });
  const writer = RecordBatchStreamWriter.writeAll(table);
  return Buffer.from(await writer.toUint8Array());
}

// `Type.Utf8View` / `Type.BinaryView` as numbers: the enum members only exist
// in Arrow 21+, and this module compiles against every supported release.
const UTF8_VIEW_TYPE_ID = 24;
const BINARY_VIEW_TYPE_ID = 23;

/**
 * Copy a Utf8View / BinaryView `Data` into a single `Data` of `type`.
 */
function materializeViewData(child: ArrowData, type: DataType): ArrowData {
  const builder = makeBuilder({ type, nullValues: [null] });
  for (const value of new Vector([child])) {
    builder.append(value);
  }
  return builder.finish().flush();
}

/**
 * Rebuild any top-level Utf8View / BinaryView column as Utf8 / Binary.
 *
 * Lance stores the view types as their offset-based equivalents anyway, so
 * nothing is lost. Doing it here also sidesteps an Arrow JS 21 bug: its IPC
 * writer emits a truncated views buffer for a *sliced* view array, which the
 * Rust reader rejects with "Need at least N bytes in buffers[0]".
 *
 * The record batches are rebuilt positionally rather than through a
 * `Record<string, Vector>`: JavaScript enumerates integer-like keys first, so
 * a field named e.g. `"1"` would otherwise be paired with the wrong column.
 *
 * Tables without view columns are returned as-is.
 */
function materializeViewColumns(table: ArrowTable): ArrowTable {
  const replacements = new Map<number, DataType>();
  table.schema.fields.forEach((field, i) => {
    if (field.type.typeId === UTF8_VIEW_TYPE_ID) {
      replacements.set(i, new Utf8());
    } else if (field.type.typeId === BINARY_VIEW_TYPE_ID) {
      replacements.set(i, new Binary());
    }
  });
  if (replacements.size === 0) {
    return table;
  }
  const fields = table.schema.fields.map((field, i) => {
    const type = replacements.get(i);
    return type === undefined
      ? field
      : new Field(field.name, type, field.nullable, field.metadata);
  });
  const schema = new Schema(fields, table.schema.metadata);
  const batches = table.batches.map((batch) => {
    const children = batch.data.children.map((child, i) => {
      const type = replacements.get(i);
      return type === undefined ? child : materializeViewData(child, type);
    });
    const data = makeData({
      type: new Struct(fields),
      length: batch.numRows,
      nullCount: 0,
      children,
    });
    return new RecordBatch(schema, data);
  });
  return new ArrowTable(schema, batches);
}

/**
 * Serialize an Arrow Table into a buffer using the Arrow IPC File serialization
 *
 * This function will apply `embeddings` to the table in a manner similar to
 * `convertToTable`.
 *
 * `schema` is required if the table is empty
 */
export async function fromTableToBuffer(
  table: ArrowTable,
  embeddings?: EmbeddingFunctionConfig,
  schema?: SchemaLike,
): Promise<Buffer> {
  if (schema !== undefined && schema !== null) {
    schema = sanitizeSchema(schema);
  }
  const tableWithEmbeddings = materializeViewColumns(
    await applyEmbeddings(table, embeddings, schema),
  );
  validateBlobSchema(tableWithEmbeddings.schema);
  const writer = RecordBatchFileWriter.writeAll(tableWithEmbeddings);
  return Buffer.from(await writer.toUint8Array());
}

/**
 * Serialize an Arrow Table into a buffer using the Arrow IPC File serialization
 *
 * This function will apply `embeddings` to the table in a manner similar to
 * `convertToTable`.
 *
 * `schema` is required if the table is empty
 */
export async function fromDataToBuffer(
  data: Data,
  embeddings?: EmbeddingFunctionConfig,
  schema?: Schema,
): Promise<Buffer> {
  if (schema !== undefined && schema !== null) {
    schema = sanitizeSchema(schema);
  }
  if (isArrowTable(data)) {
    const table = sanitizeTable(data);
    // If we have a schema with embedding functions, we need to ensure all columns exist
    // before applying embeddings, since applyEmbeddingsFromMetadata expects all columns
    // to be present in the table
    if (schema && schema.metadata?.has("embedding_functions")) {
      const alignedTable = alignTableToSchema(table, schema);
      return fromTableToBuffer(alignedTable, embeddings, schema);
    } else {
      return fromTableToBuffer(table, embeddings, schema);
    }
  } else {
    const table = await convertToTable(data, embeddings, { schema });
    return fromTableToBuffer(table);
  }
}

/**
 * Read a single record batch from a buffer.
 *
 * Returns null if the buffer does not contain a record batch
 */
export async function fromBufferToRecordBatch(
  data: Buffer,
): Promise<RecordBatch | null> {
  const iter = await RecordBatchFileReader.readAll(Buffer.from(data)).next()
    .value;
  const recordBatch = iter?.next().value;
  return recordBatch || null;
}

/**
 * Create a buffer containing a single record batch
 */
export async function fromRecordBatchToBuffer(
  batch: RecordBatch,
): Promise<Buffer> {
  const writer = new RecordBatchFileWriter().writeAll([batch]);
  return Buffer.from(await writer.toUint8Array());
}

/**
 * Create a buffer containing a single record batch using the Arrow IPC Stream
 * serialization. Each call produces a self-contained Stream message (schema +
 * batch + EOS) suitable for incremental decode by `arrow_ipc::reader::StreamReader`.
 */
export async function fromRecordBatchToStreamBuffer(
  batch: RecordBatch,
): Promise<Buffer> {
  const table = materializeViewColumns(new ArrowTable([batch]));
  const writer = RecordBatchStreamWriter.writeAll(table);
  return Buffer.from(await writer.toUint8Array());
}

/**
 * Serialize an Arrow Table into a buffer using the Arrow IPC Stream serialization
 *
 * This function will apply `embeddings` to the table in a manner similar to
 * `convertToTable`.
 *
 * `schema` is required if the table is empty
 */
export async function fromTableToStreamBuffer(
  table: ArrowTable,
  embeddings?: EmbeddingFunctionConfig,
  schema?: SchemaLike,
): Promise<Buffer> {
  const tableWithEmbeddings = materializeViewColumns(
    await applyEmbeddings(table, embeddings, schema),
  );
  const writer = RecordBatchStreamWriter.writeAll(tableWithEmbeddings);
  return Buffer.from(await writer.toUint8Array());
}

/**
 * Reorder the columns in `batch` so that they agree with the field order in `schema`
 */
function alignBatch(batch: RecordBatch, schema: Schema): RecordBatch {
  const alignedChildren = [];
  for (const field of schema.fields) {
    const indexInBatch = batch.schema.fields?.findIndex(
      (f) => f.name === field.name,
    );
    if (indexInBatch < 0) {
      throw new Error(
        `The column ${field.name} was not found in the Arrow Table`,
      );
    }
    alignedChildren.push(batch.data.children[indexInBatch]);
  }
  const newData = makeData({
    type: new Struct(schema.fields),
    length: batch.numRows,
    nullCount: batch.nullCount,
    children: alignedChildren as unknown as ArrowData<DataType>[],
  });
  return new RecordBatch(schema, newData);
}

/**
 * Reorder the columns in `table` so that they agree with the field order in `schema`
 */
function alignTable(table: ArrowTable, schema: Schema): ArrowTable {
  const alignedBatches = table.batches.map((batch) =>
    alignBatch(batch, schema),
  );
  return new ArrowTable(schema, alignedBatches);
}

/**
 * Create an empty table with the given schema
 */
export function createEmptyTable(schema: Schema): ArrowTable {
  return new ArrowTable(sanitizeSchema(schema));
}

function validateSchemaEmbeddings(
  schema: Schema,
  data: Array<Record<string, unknown>>,
  embeddings: EmbeddingFunctionConfig | undefined,
): Schema {
  const fields = [];
  const missingEmbeddingFields = [];

  // First we check if the field is a `FixedSizeList`
  // Then we check if the data contains the field
  // if it does not, we add it to the list of missing embedding fields
  // Finally, we check if those missing embedding fields are `this._embeddings`
  // if they are not, we throw an error
  for (let field of schema.fields) {
    if (isFixedSizeList(field.type)) {
      field = sanitizeField(field);
      if (data.length !== 0 && data?.[0]?.[field.name] === undefined) {
        // Check if there's an embedding function registered for this field
        let hasEmbeddingFunction = false;

        // Check schema metadata for embedding functions
        if (schema.metadata.has("embedding_functions")) {
          const entries = parseEmbeddingMetadata(
            schema.metadata.get("embedding_functions")!,
          );
          if (entries.some((f) => f.vectorColumn === field.name)) {
            hasEmbeddingFunction = true;
          }
        }

        // Check passed embedding function parameter
        if (embeddings && embeddings.vectorColumn === field.name) {
          hasEmbeddingFunction = true;
        }

        // If the field is nullable AND there's no embedding function, allow undefined/omitted values
        if (field.nullable && !hasEmbeddingFunction) {
          fields.push(field);
        } else {
          // Either not nullable OR has embedding function - require explicit values
          if (hasEmbeddingFunction) {
            // Don't add to missingEmbeddingFields since this is expected to be filled by embedding function
            fields.push(field);
          } else {
            missingEmbeddingFields.push(field);
          }
        }
      } else {
        fields.push(field);
      }
    } else {
      fields.push(field);
    }
  }

  if (missingEmbeddingFields.length > 0 && embeddings === undefined) {
    throw new Error(
      `Table has embeddings: "${missingEmbeddingFields
        .map((f) => f.name)
        .join(",")}", but no embedding function was provided`,
    );
  }

  return new Schema(fields, schema.metadata);
}

/**
 * Ensures that all nested fields defined in the schema exist in the data,
 * filling missing fields with null values.
 */
export function ensureNestedFieldsExist(
  data: Array<Record<string, unknown>>,
  schema: Schema,
): Array<Record<string, unknown>> {
  return data.map((row) => {
    const completeRow: Record<string, unknown> = {};

    for (const field of schema.fields) {
      if (field.name in row) {
        if (
          isPlainStructField(field) &&
          row[field.name] !== null &&
          row[field.name] !== undefined
        ) {
          // Handle nested struct
          const nestedValue = row[field.name] as Record<string, unknown>;
          completeRow[field.name] = ensureStructFieldsExist(
            nestedValue,
            field.type,
          );
        } else {
          // Non-struct field or null struct value
          completeRow[field.name] = row[field.name];
        }
      } else {
        // Keep a missing struct valid while filling each of its children with
        // null. This is distinct from an explicitly null struct value.
        completeRow[field.name] = isPlainStructField(field)
          ? ensureStructFieldsExist({}, field.type as Struct)
          : null;
      }
    }

    return completeRow;
  });
}

/**
 * Blob fields are Arrow structs, but their values are bytes, URIs, or
 * `{ data } | { uri }` inputs that the blob coercion in `makeArrowTable`
 * handles, so they must not be filled in like ordinary structs.
 */
function isPlainStructField(field: Field): boolean {
  return field.type.constructor.name === "Struct" && !isBlobField(field);
}

/**
 * Recursively ensures that all fields in a struct type exist in the data,
 * filling missing fields with null values.
 */
function ensureStructFieldsExist(
  data: Record<string, unknown>,
  structType: Struct,
): Record<string, unknown> {
  const completeStruct: Record<string, unknown> = {};

  for (const childField of structType.children) {
    if (childField.name in data) {
      if (
        isPlainStructField(childField) &&
        data[childField.name] !== null &&
        data[childField.name] !== undefined
      ) {
        // Recursively handle nested struct
        completeStruct[childField.name] = ensureStructFieldsExist(
          data[childField.name] as Record<string, unknown>,
          childField.type,
        );
      } else {
        // Non-struct field or null struct value
        completeStruct[childField.name] = data[childField.name];
      }
    } else {
      // Keep a missing struct valid while filling each of its children with
      // null. This is distinct from an explicitly null struct value.
      completeStruct[childField.name] = isPlainStructField(childField)
        ? ensureStructFieldsExist({}, childField.type as Struct)
        : null;
    }
  }

  return completeStruct;
}

interface JsonDataType {
  type: string;
  fields?: JsonField[];
  length?: number;
}

interface JsonField {
  name: string;
  type: JsonDataType;
  nullable: boolean;
  metadata: Map<string, string>;
}

// Matches format of https://github.com/lancedb/lance/blob/main/rust/lance/src/arrow/json.rs
export function dataTypeToJson(dataType: DataType): JsonDataType {
  switch (dataType.typeId) {
    // For primitives, matches https://github.com/lancedb/lance/blob/e12bb9eff2a52f753668d4b62c52e4d72b10d294/rust/lance-core/src/datatypes.rs#L185
    case Type.Null:
      return { type: "null" };
    case Type.Bool:
      return { type: "bool" };
    case Type.Int8:
      return { type: "int8" };
    case Type.Int16:
      return { type: "int16" };
    case Type.Int32:
      return { type: "int32" };
    case Type.Int64:
      return { type: "int64" };
    case Type.Uint8:
      return { type: "uint8" };
    case Type.Uint16:
      return { type: "uint16" };
    case Type.Uint32:
      return { type: "uint32" };
    case Type.Uint64:
      return { type: "uint64" };
    case Type.Int: {
      const bitWidth = (dataType as Int).bitWidth;
      const signed = (dataType as Int).isSigned;
      const prefix = signed ? "" : "u";
      return { type: `${prefix}int${bitWidth}` };
    }
    case Type.Float: {
      switch ((dataType as Float).precision) {
        case Precision.HALF:
          return { type: "halffloat" };
        case Precision.SINGLE:
          return { type: "float" };
        case Precision.DOUBLE:
          return { type: "double" };
      }
      throw Error("Unsupported float precision");
    }
    case Type.Float16:
      return { type: "halffloat" };
    case Type.Float32:
      return { type: "float" };
    case Type.Float64:
      return { type: "double" };
    case Type.Utf8:
      return { type: "string" };
    case Type.Binary:
      return { type: "binary" };
    case Type.LargeUtf8:
      return { type: "large_string" };
    case Type.LargeBinary:
      return { type: "large_binary" };
    case Type.List:
      return {
        type: "list",
        fields: [fieldToJson((dataType as List).children[0])],
      };
    case Type.FixedSizeList: {
      const fixedSizeList = dataType as FixedSizeList;
      return {
        type: "fixed_size_list",
        fields: [fieldToJson(fixedSizeList.children[0])],
        length: fixedSizeList.listSize,
      };
    }
    case Type.Struct:
      return {
        type: "struct",
        fields: (dataType as Struct).children.map(fieldToJson),
      };
    case Type.Date: {
      const unit = (dataType as Date_).unit;
      return {
        type: unit === DateUnit.DAY ? "date32:day" : "date64:ms",
      };
    }
    case Type.Timestamp: {
      const timestamp = dataType as Timestamp;
      const timezone = timestamp.timezone || "-";
      return {
        type: `timestamp:${timestamp.unit}:${timezone}`,
      };
    }
    case Type.Decimal: {
      const decimal = dataType as Decimal;
      return {
        type: `decimal:${decimal.bitWidth}:${decimal.precision}:${decimal.scale}`,
      };
    }
    case Type.Duration: {
      const duration = dataType as Duration;
      return { type: `duration:${duration.unit}` };
    }
    case Type.FixedSizeBinary: {
      const byteWidth = (dataType as FixedSizeBinary).byteWidth;
      return { type: `fixed_size_binary:${byteWidth}` };
    }
    case Type.Dictionary: {
      const dict = dataType as Dictionary;
      const indexType = dataTypeToJson(dict.indices);
      const valueType = dataTypeToJson(dict.valueType);
      return {
        type: `dict:${valueType.type}:${indexType.type}:false`,
      };
    }
  }
  throw new Error("Unsupported data type");
}

function fieldToJson(field: Field): JsonField {
  return {
    name: field.name,
    type: dataTypeToJson(field.type),
    nullable: field.nullable,
    metadata: field.metadata,
  };
}

function alignTableToSchema(
  table: ArrowTable,
  targetSchema: Schema,
): ArrowTable {
  const existingColumns = new Map<string, Vector>();

  // Map existing columns
  for (const field of table.schema.fields) {
    existingColumns.set(field.name, table.getChild(field.name)!);
  }

  // Create vectors for all fields in target schema
  const alignedColumns: Record<string, Vector> = {};

  for (const field of targetSchema.fields) {
    if (existingColumns.has(field.name)) {
      // Column exists, use it
      alignedColumns[field.name] = existingColumns.get(field.name)!;
    } else {
      // Column missing, create null vector
      alignedColumns[field.name] = createNullVector(field, table.numRows);
    }
  }

  // Create new table with aligned schema and columns
  return new ArrowTable(targetSchema, alignedColumns);
}

function createNullVector(field: Field, numRows: number): Vector {
  if (field.type.constructor.name === "Struct") {
    // For struct types, create a struct with null fields
    const structType = field.type as Struct;
    const childVectors = structType.children.map((childField) =>
      createNullVector(childField, numRows),
    );

    // Create struct data
    const structData = makeData({
      type: structType,
      length: numRows,
      nullCount: 0,
      children: childVectors.map((v) => v.data[0]),
    });

    return arrowMakeVector(structData);
  } else {
    // For other types, create a vector of nulls
    const nullBitmap = new Uint8Array(Math.ceil(numRows / 8));
    // All bits are 0, meaning all values are null

    const data = makeData({
      type: field.type,
      length: numRows,
      nullCount: numRows,
      nullBitmap,
    });

    return arrowMakeVector(data);
  }
}