-- 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 testify = testimony.new("== mneme: ft search ==")

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

-- Helper: hand-compute BM25 for verification
local bm25_score = function(tf, df, N, doc_len, avgdl, k1, b)
	local idf = math.log(((N - df + 0.5) / (df + 0.5)) + 1.0)
	local num = tf * (k1 + 1)
	local denom = tf + k1 * (1 - b + b * doc_len / avgdl)
	return idf * num / denom
end

testify:that("single-term query returns ranked results", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	ft:put("d1", "the cat sat on the mat")
	ft:put("d2", "the dog sat on the log")
	ft:put("d3", "the cat chased the dog around the yard")

	local results = ft:search("cat")
	testimony.assert_true(#results > 0)

	-- Both d1 and d3 contain "cat", d2 does not
	local keys = {}
	for _, r in ipairs(results) do
		keys[r.key] = true
	end
	testimony.assert_true(keys["d1"] ~= nil)
	testimony.assert_true(keys["d3"] ~= nil)
	testimony.assert_nil(keys["d2"])

	-- Scores should be positive
	for _, r in ipairs(results) do
		testimony.assert_true(r.score > 0)
	end

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

testify:that("multi-term query accumulates scores", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	ft:put("d1", "congestion window flow control")
	ft:put("d2", "congestion control algorithm")
	ft:put("d3", "flow rate measurement")

	-- Query with two terms that both appear in d1
	local results = ft:search("congestion flow")
	testimony.assert_true(#results > 0)

	-- d1 should rank highest (matches both terms)
	testimony.assert_equal("d1", results[1].key)

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

testify:that("BM25 scores match hand-computed values", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	-- Carefully controlled corpus for exact verification
	-- Using default analyzer: min 2 chars, lowercased
	ft:put("d1", "aa bb cc") -- 3 terms: aa, bb, cc
	ft:put("d2", "aa aa bb dd") -- 4 terms: aa(x2), bb, dd

	local N = 2
	local avgdl = (3 + 4) / 2 -- 3.5
	local k1 = 1.2
	local b = 0.75

	-- Search for "aa": df=2 (both docs), tf(d1)=1, tf(d2)=2
	local results = ft:search("aa")
	testimony.assert_equal(2, #results)

	local expected_d1 = bm25_score(1, 2, N, 3, avgdl, k1, b)
	local expected_d2 = bm25_score(2, 2, N, 4, avgdl, k1, b)

	-- Find each result
	local actual = {}
	for _, r in ipairs(results) do
		actual[r.key] = r.score
	end

	testimony.assert_true(
		math.abs(actual["d1"] - expected_d1) < 0.01,
		"d1 score mismatch: got " .. tostring(actual["d1"]) .. " expected " .. tostring(expected_d1)
	)
	testimony.assert_true(
		math.abs(actual["d2"] - expected_d2) < 0.01,
		"d2 score mismatch: got " .. tostring(actual["d2"]) .. " expected " .. tostring(expected_d2)
	)

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

testify:that("empty query returns empty results", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	ft:put("d1", "hello world")

	local results = ft:search("")
	testimony.assert_equal(0, #results)

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

testify:that("query with no matches returns empty", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	ft:put("d1", "hello world")

	local results = ft:search("xyzzy notaword")
	testimony.assert_equal(0, #results)

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

testify:that("top_k limits result count", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	for i = 1, 20 do
		ft:put("d" .. i, "common term here doc " .. i)
	end

	local results = ft:search("common", { top_k = 5 })
	testimony.assert_true(#results <= 5)
	testimony.assert_true(#results > 0)

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

testify:that("results are sorted by descending score", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	ft:put("d1", "target target target") -- tf=3
	ft:put("d2", "target something else") -- tf=1
	ft:put("d3", "target target here now") -- tf=2
	ft:put("d4", "nothing related at all") -- tf=0

	local results = ft:search("target")
	testimony.assert_true(#results >= 2)

	-- Verify descending order
	for i = 2, #results do
		testimony.assert_true(
			results[i - 1].score >= results[i].score,
			"results not sorted descending at position " .. i
		)
	end

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

testify:that("search on empty corpus returns empty", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	local results = ft:search("hello")
	testimony.assert_equal(0, #results)

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

testify:that("custom k1 and b parameters affect scoring", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	ft:put("d1", "word word word word word") -- high tf
	ft:put("d2", "word other stuff here now") -- lower tf

	local r1 = ft:search("word", { k1 = 0.0 }) -- k1=0 → tf doesn't matter
	local r2 = ft:search("word", { k1 = 100.0 }) -- k1=100 → tf matters a lot

	-- With k1=0, both should have similar scores (tf saturated)
	-- With k1=100, d1 should dominate due to higher tf
	testimony.assert_true(#r1 >= 2)
	testimony.assert_true(#r2 >= 2)

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

testify:that("search after delete reflects updated corpus", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	ft:put("d1", "hello world test")
	ft:put("d2", "hello again here")

	local r1 = ft:search("hello")
	testimony.assert_equal(2, #r1)

	ft:del("d1")

	local r2 = ft:search("hello")
	testimony.assert_equal(1, #r2)
	testimony.assert_equal("d2", r2[1].key)

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

testify:that("search after update reflects new content", function()
	local path = tmp_path()
	local db = mneme.open(path, { sync = "none" })
	local ft = db:ft_keyspace("corpus")

	ft:put("d1", "alpha beta gamma")
	ft:put("d2", "delta epsilon zeta")

	local r1 = ft:search("alpha")
	testimony.assert_equal(1, #r1)
	testimony.assert_equal("d1", r1[1].key)

	-- Update d1 to no longer contain "alpha"
	ft:put("d1", "delta epsilon new content")

	local r2 = ft:search("alpha")
	testimony.assert_equal(0, #r2)

	-- Now d1 should match "delta"
	local r3 = ft:search("delta")
	testimony.assert_equal(2, #r3)

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

testify:conclude()