-- SPDX-FileCopyrightText: © 2026 Vladimir Zorin <vladimir@deviant.guru>
-- SPDX-License-Identifier: OWL-1.0 or later

local testimony = require("testimony")
local std = require("std")
local mneme = require("mneme")
local core = require("mneme.core")

local testify = testimony.new("== mneme: vectors ==")

local tmp_path = function()
	return "/tmp/mneme_vec_" .. std.nanoid() .. ".mneme"
end

-- Helper: generate a random-ish unit vector of given dimension
local make_vec = function(dim, seed)
	local v = {}
	local s = seed or 1
	for i = 1, dim do
		s = (s * 1103515245 + 12345) % (2 ^ 31)
		v[i] = (s / (2 ^ 31)) * 2 - 1
	end
	-- Normalize
	local mag = 0
	for i = 1, dim do
		mag = mag + v[i] * v[i]
	end
	mag = math.sqrt(mag)
	for i = 1, dim do
		v[i] = v[i] / mag
	end
	return v
end

testify:that("vec_dot computes correct dot product", function()
	local a = core.pack_floats({ 1, 0, 0, 0 })
	local b = core.pack_floats({ 0, 1, 0, 0 })
	local c = core.pack_floats({ 1, 0, 0, 0 })

	-- Orthogonal vectors: dot = 0
	-- For now test via unpack round-trip
	local at = core.unpack_floats(a)
	testimony.assert_equal(4, #at)
	testimony.assert_true(math.abs(at[1] - 1.0) < 0.001)
	testimony.assert_true(math.abs(at[2] - 0.0) < 0.001)
end)

testify:that("vset and vget round-trip a vector", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ks = db:keyspace("vecs")

	local vec = { 0.5, -0.3, 0.8, 0.1 }
	local ok, err = ks:vset("v1", vec)
	testimony.assert_not_nil(ok, err)

	local got = ks:vget("v1")
	testimony.assert_not_nil(got)
	testimony.assert_equal(4, #got)
	-- Float32 precision: ~7 decimal digits
	for i = 1, 4 do
		testimony.assert_true(math.abs(got[i] - vec[i]) < 0.0001)
	end

	db:close()
	std.fs.remove(path)
end)

testify:that("vdim returns stored dimension", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ks = db:keyspace("vecs")

	testimony.assert_nil(ks:vdim()) -- no vectors yet

	ks:vset("v1", { 1, 2, 3, 4, 5, 6, 7, 8 })
	testimony.assert_equal(8, ks:vdim())

	db:close()
	std.fs.remove(path)
end)

testify:that("vset dimension mismatch returns error", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ks = db:keyspace("vecs")

	ks:vset("v1", { 1, 2, 3, 4 }) -- dim=4

	local ok, err = ks:vset("v2", { 1, 2, 3, 4, 5, 6, 7, 8 }) -- dim=8
	testimony.assert_nil(ok)
	testimony.assert_equal("dimension mismatch", err)

	db:close()
	std.fs.remove(path)
end)

testify:that("vsearch cosine finds nearest vectors", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ks = db:keyspace("vecs")

	local dim = 16
	-- Insert 100 vectors
	for i = 1, 100 do
		ks:vset("doc:" .. string.format("%03d", i), make_vec(dim, i))
	end

	-- Query with a vector very close to doc:042
	local query = make_vec(dim, 42)
	-- Perturb slightly
	query[1] = query[1] + 0.01

	local results = ks:vsearch(query, {
		metric = "cosine",
		top_k = 5,
	})

	testimony.assert_equal(5, #results)
	-- The closest should be doc:042 (cosine distance ≈ 0)
	testimony.assert_equal("doc:042", results[1].key)
	testimony.assert_true(results[1].score < 0.01)

	-- Results should be sorted by score (ascending for cosine)
	for i = 2, #results do
		testimony.assert_true(results[i].score >= results[i - 1].score)
	end

	db:close()
	std.fs.remove(path)
end)

testify:that("vsearch with prefix filter", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ks = db:keyspace("vecs")

	local dim = 8
	ks:vset("cat:1", make_vec(dim, 10))
	ks:vset("cat:2", make_vec(dim, 20))
	ks:vset("dog:1", make_vec(dim, 30))
	ks:vset("dog:2", make_vec(dim, 40))

	local results = ks:vsearch(make_vec(dim, 10), {
		metric = "cosine",
		top_k = 10,
		filter = "cat:",
	})

	-- Should only return cat: entries
	testimony.assert_equal(2, #results)
	for _, r in ipairs(results) do
		testimony.assert_true(r.key:sub(1, 4) == "cat:")
	end

	db:close()
	std.fs.remove(path)
end)

testify:that("vsearch dot product metric", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ks = db:keyspace("vecs")

	local dim = 4
	ks:vset("a", { 1, 0, 0, 0 })
	ks:vset("b", { 0, 1, 0, 0 })
	ks:vset("c", { 0.9, 0.1, 0, 0 })

	local results = ks:vsearch({ 1, 0, 0, 0 }, {
		metric = "dot",
		top_k = 3,
	})

	testimony.assert_equal(3, #results)
	-- For dot product, higher is better. Results sorted descending.
	-- "a" should have highest dot product (1.0), then "c" (~0.9), then "b" (0.0)
	testimony.assert_equal("a", results[1].key)
	testimony.assert_true(math.abs(results[1].score - 1.0) < 0.01)

	db:close()
	std.fs.remove(path)
end)

testify:that("vsearch l2 metric", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ks = db:keyspace("vecs")

	ks:vset("near", { 1.0, 0.0, 0.0, 0.0 })
	ks:vset("far", { -1.0, 0.0, 0.0, 0.0 })

	local results = ks:vsearch({ 0.9, 0.0, 0.0, 0.0 }, {
		metric = "l2",
		top_k = 2,
	})

	testimony.assert_equal(2, #results)
	testimony.assert_equal("near", results[1].key) -- closer
	testimony.assert_true(results[1].score < results[2].score)

	db:close()
	std.fs.remove(path)
end)

testify:that("vsearch on empty keyspace returns empty table", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ks = db:keyspace("empty")

	local results = ks:vsearch({ 1, 0, 0 }, { metric = "cosine", top_k = 5 })
	testimony.assert_equal(0, #results)

	db:close()
	std.fs.remove(path)
end)

testify:that("vsearch finds overflow vectors (high dimension)", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ks = db:keyspace("bigvecs")

	local dim = 512 -- 2 + 512*4 = 2050 bytes, exceeds inline limit
	-- Insert 20 vectors
	for i = 1, 20 do
		ks:vset("v:" .. string.format("%03d", i), make_vec(dim, i))
	end

	-- Query with a vector very close to v:007
	local query = make_vec(dim, 7)
	query[1] = query[1] + 0.01

	local results = ks:vsearch(query, {
		metric = "cosine",
		top_k = 5,
	})

	testimony.assert_equal(5, #results)
	-- The closest should be v:007
	testimony.assert_equal("v:007", results[1].key)
	testimony.assert_true(results[1].score < 0.01)

	-- Results should be sorted by score (ascending for cosine)
	for i = 2, #results do
		testimony.assert_true(results[i].score >= results[i - 1].score)
	end

	db:close()
	std.fs.remove(path)
end)

testify:that("vec_normalize produces unit vector", function()
	local v = { 3, 4 }
	-- Manually normalize: magnitude = 5, normalized = (0.6, 0.8)
	local packed = core.pack_floats(v)
	testimony.assert_not_nil(packed)
	local unpacked = core.unpack_floats(packed)
	testimony.assert_equal(2, #unpacked)
	testimony.assert_true(math.abs(unpacked[1] - 3.0) < 0.001)
	testimony.assert_true(math.abs(unpacked[2] - 4.0) < 0.001)
end)

testify:conclude()