// SPDX-FileCopyrightText: 2025-2026 Vasyl Zakharchenko
// SPDX-License-Identifier: MIT

import { sql, type SQL, type AnyColumn } from "drizzle-orm";
export { vectorTiDBType } from "../customTypes";

function validateVectorValue(value: number[]): void {
  for (const item of value) {
    if (typeof item !== "number" || !Number.isFinite(item)) {
      throw new Error("TiDB vector contains invalid number");
    }
  }
}

function vectorToText(value: number[]): string {
  validateVectorValue(value);
  return `[${value.join(",")}]`;
}

/**
 * Input accepted by TiDB vector SQL helpers.
 *
 * - `number[]`: converted to `VEC_FROM_TEXT('[...]')`
 * - `string`: treated as raw SQL vector expression, for example:
 *   `CAST('[0.3, 0.5, -0.1]' AS VECTOR)`
 * - `SQL` / `AnyColumn`: passed through as-is
 */
type VectorInput = string | number[] | SQL | AnyColumn;

function vectorExpr(value: VectorInput): SQL {
  if (Array.isArray(value)) {
    return sql`VEC_FROM_TEXT(${vectorToText(value)})`;
  }
  if (typeof value === "string") {
    return sql.raw(value);
  }
  return sql`${value}`;
}

/**
 * Converts a text representation of a vector into a TiDB VECTOR expression.
 *
 * TiDB function: `VEC_FROM_TEXT(string)`.
 *
 * @example
 * vecFromText("[1, 2, 3]")
 */
export function vecFromText(text: string): SQL {
  return sql`VEC_FROM_TEXT(${text})`;
}

/**
 * Converts a vector value/expression to its normalized string representation.
 *
 * TiDB function: `VEC_AS_TEXT(vector)`.
 *
 * @example
 * vecAsText(table.embedding)
 */
export function vecAsText(vector: VectorInput): SQL<string> {
  return sql<string>`VEC_AS_TEXT(${vectorExpr(vector)})`;
}

/**
 * Returns the dimension (number of elements) of a vector.
 *
 * TiDB function: `VEC_DIMS(vector)`.
 *
 * Accepts vector column/expression, `sql.raw(...)`, raw SQL string expression,
 * or `number[]`.
 *
 * @example
 * vecDims(table.embedding)
 * vecDims(sql.raw("CAST('[0.3, 0.5, -0.1]' AS VECTOR)"))
 * vecDims("CAST('[0.3, 0.5, -0.1]' AS VECTOR)")
 */
export function vecDims(vector: VectorInput): SQL<number> {
  return sql<number>`VEC_DIMS(${vectorExpr(vector)})`;
}

/**
 * Calculates L2 norm (Euclidean norm) of a vector.
 *
 * TiDB function: `VEC_L2_NORM(vector)`.
 */
export function vecL2Norm(vector: VectorInput): SQL<number> {
  return sql`VEC_L2_NORM(${vectorExpr(vector)})`;
}

/**
 * Calculates L2 distance (Euclidean distance) between two vectors.
 *
 * TiDB function: `VEC_L2_DISTANCE(vector1, vector2)`.
 *
 * Both vectors must have the same dimensions.
 */
export function vecL2Distance(left: VectorInput, right: VectorInput): SQL<number> {
  return sql<number>`VEC_L2_DISTANCE(${vectorExpr(left)}, ${vectorExpr(right)})`;
}

/**
 * Calculates cosine distance between two vectors.
 *
 * TiDB function: `VEC_COSINE_DISTANCE(vector1, vector2)`.
 *
 * Both vectors must have the same dimensions.
 */
export function vecCosineDistance(left: VectorInput, right: VectorInput): SQL<number> {
  return sql<number>`VEC_COSINE_DISTANCE(${vectorExpr(left)}, ${vectorExpr(right)})`;
}

/**
 * Calculates negative inner product distance between two vectors.
 *
 * TiDB function: `VEC_NEGATIVE_INNER_PRODUCT(vector1, vector2)`.
 *
 * Both vectors must have the same dimensions.
 */
export function vecNegativeInnerProduct(left: VectorInput, right: VectorInput): SQL<number> {
  return sql<number>`VEC_NEGATIVE_INNER_PRODUCT(${vectorExpr(left)}, ${vectorExpr(right)})`;
}

/**
 * Calculates L1 distance (Manhattan distance) between two vectors.
 *
 * TiDB function: `VEC_L1_DISTANCE(vector1, vector2)`.
 *
 * Both vectors must have the same dimensions.
 */
export function vecL1Distance(left: VectorInput, right: VectorInput): SQL<number> {
  return sql<number>`VEC_L1_DISTANCE(${vectorExpr(left)}, ${vectorExpr(right)})`;
}