sift/internal/search/engine_test.go

package search

import (
	"context"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/dcache"
)

type fakeSrc struct {
	name    string
	hits    []Hit
	warn    string
	delay   time.Duration
	credits int
}

func (f fakeSrc) Name() string { return f.name }
func (f fakeSrc) Search(ctx context.Context, query string, k int) SourceOut {
	if f.delay > 0 {
		select {
		case <-time.After(f.delay):
		case <-ctx.Done():
			return SourceOut{Warn: f.name + ": timeout"}
		}
	}
	return SourceOut{Hits: f.hits, Warn: f.warn, Credits: f.credits}
}

type stubEmb struct {
	ok     bool
	reason string
}

func (s stubEmb) Embed(ctx context.Context, inputs []string) ([][]float64, error) {
	return nil, context.Canceled
}
func (s stubEmb) Acquire(ctx context.Context) (bool, string) { return s.ok, s.reason }
func (s stubEmb) Release()                                   {}

func TestFanInPartial(t *testing.T) {
	e := &Engine{
		Sources: []Source{
			fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Title: "ok", Snippet: "x", Source: NameSearxNG}}},
			fakeSrc{name: NameBraveLLM, warn: "brave_llm: 500"},
			fakeSrc{name: NameTavily, delay: 2 * time.Second},
		},
		Timeout:  50 * time.Millisecond,
		DefaultK: 8,
		RerankOn: false,
	}
	r := e.Run(context.Background(), "q", 8, false, "titan", dcache.ClassFleet)
	if len(r.Hits) != 1 || r.Hits[0].URL != "https://ok.example/a" {
		t.Fatalf("hits %+v", r.Hits)
	}
	joined := strings.Join(r.Warnings, " ")
	if !strings.Contains(joined, "500") || !strings.Contains(joined, "timeout") {
		t.Fatalf("warnings %v", r.Warnings)
	}
}

func TestCacheClassAndHour(t *testing.T) {
	db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	now := time.Date(2026, 9, 23, 10, 0, 0, 0, time.UTC)
	src := fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Title: "ok", Snippet: "x", Source: NameSearxNG}}}
	e := &Engine{Sources: []Source{src}, DefaultK: 8, Cache: db, RerankOn: false, Now: func() time.Time { return now }}
	r1 := e.Run(context.Background(), "Hello World", 8, false, "titan", dcache.ClassFleet)
	if r1.Cached {
		t.Fatal("first miss")
	}
	r2 := e.Run(context.Background(), "hello   world", 8, false, "titan", dcache.ClassFleet)
	if !r2.Cached {
		t.Fatal("norm/hour hit")
	}
	r3 := e.Run(context.Background(), "hello world", 8, false, "gs", dcache.ClassRestricted)
	if r3.Cached {
		t.Fatal("restricted must miss fleet cache")
	}
	e.Now = func() time.Time { return now.Add(time.Hour) }
	r4 := e.Run(context.Background(), "hello world", 8, false, "titan", dcache.ClassFleet)
	if r4.Cached {
		t.Fatal("next hour must miss")
	}
}

func TestNoCacheOnDegraded(t *testing.T) {
	db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	e := &Engine{
		Sources: []Source{
			fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Source: NameSearxNG}}},
			fakeSrc{name: NameBraveLLM, warn: "brave_llm: 500"},
		},
		DefaultK: 8, Cache: db, RerankOn: false,
	}
	e.Run(context.Background(), "q", 8, false, "titan", dcache.ClassFleet)
	r := e.Run(context.Background(), "q", 8, false, "titan", dcache.ClassFleet)
	if r.Cached {
		t.Fatal("degraded cached")
	}
}

func TestRerankFallbackStill200(t *testing.T) {
	e := &Engine{
		Sources:  []Source{fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Source: NameSearxNG}}}},
		DefaultK: 8,
		RerankOn: true,
		Embedder: stubEmb{ok: false, reason: "queue_full"},
	}
	r := e.Run(context.Background(), "q", 8, true, "titan", dcache.ClassFleet)
	if len(r.Hits) != 1 {
		t.Fatal(r.Hits)
	}
	if strings.Contains(strings.Join(r.Warnings, " "), "rerank fell back") {
		t.Fatalf("RRF must not depend on embed: %v", r.Warnings)
	}
}

func TestRRFInterleaveAndNoRerank(t *testing.T) {
	searx := []Hit{
		{URL: "https://a.example/s", Title: "s1", Source: NameSearxNG},
		{URL: "https://b.example/s", Title: "s2", Source: NameSearxNG},
	}
	tav := []Hit{
		{URL: "https://tavily.example/1", Title: "t1", Source: NameTavily},
		{URL: "https://a.example/s", Title: "t-dup", Snippet: "longer tavily snippet here", Source: NameTavily},
	}
	merged := rrfMerge([][]Hit{searx, tav})
	if len(merged) != 3 {
		t.Fatalf("%+v", merged)
	}
	if merged[0].URL != "https://a.example/s" {
		t.Fatalf("shared URL should rank first, got %+v", merged[0])
	}
	e := &Engine{
		Sources: []Source{
			fakeSrc{name: NameSearxNG, hits: searx},
			fakeSrc{name: NameTavily, hits: tav},
		},
		DefaultK: 8,
		RerankOn: true,
		NoRerank: map[string]bool{NameTavily: true},
	}
	r := e.Run(context.Background(), "q", 8, true, "titan", dcache.ClassFleet)
	if r.Hits[0].URL != "https://tavily.example/1" {
		t.Fatalf("no_rerank tavily should lead: %+v", r.Hits)
	}
}

func TestRRFKeepsOriginalURL(t *testing.T) {
	orig := "https://ex.com/docs/"
	merged := rrfMerge([][]Hit{
		{{URL: orig, Title: "slash", Source: NameSearxNG}},
		{{URL: "https://ex.com/docs", Title: "stripped", Snippet: "longer", Source: NameTavily}},
	})
	if len(merged) != 1 {
		t.Fatalf("%+v", merged)
	}
	if merged[0].URL != orig {
		t.Fatalf("want original URL %q got %q", orig, merged[0].URL)
	}
	if CanonicalURL(orig) != "https://ex.com/docs" {
		t.Fatalf("canonical %q", CanonicalURL(orig))
	}
}

func TestDailyCapSkipsTavily(t *testing.T) {
	db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	var n int
	tav := fakeSrc{name: NameTavily, hits: []Hit{{URL: "https://t.example/a", Source: NameTavily}}, credits: 1}
	e := &Engine{Sources: []Source{tav}, DefaultK: 8, Cache: db, TavilyCap: 1, RerankOn: false}
	e.Run(context.Background(), "q1", 8, false, "titan", dcache.ClassFleet)
	r := e.Run(context.Background(), "q2", 8, false, "titan", dcache.ClassFleet)
	for _, w := range r.Warnings {
		if strings.Contains(w, "daily cap") {
			n++
		}
	}
	if n != 1 {
		t.Fatalf("warnings %v", r.Warnings)
	}
}

func TestRestrictedSkipsTavilyCap(t *testing.T) {
	db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	tav := fakeSrc{name: NameTavily, hits: []Hit{{URL: "https://t.example/a", Source: NameTavily}}, credits: 1}
	e := &Engine{Sources: []Source{tav}, DefaultK: 8, Cache: db, TavilyCap: 1, RerankOn: false}
	e.Run(context.Background(), "q1", 8, false, "titan", dcache.ClassFleet)
	day := e.now().UTC().Format("2006-01-02")
	if db.GetSearchCredits(NameTavily, day) != 1 {
		t.Fatalf("fleet row %d", db.GetSearchCredits(NameTavily, day))
	}
	r := e.Run(context.Background(), "q-gs", 8, false, "straylight-gs", dcache.ClassRestricted)
	if len(r.Hits) != 1 || r.Hits[0].URL != "https://t.example/a" {
		t.Fatalf("restricted should still get Tavily: %+v %v", r.Hits, r.Warnings)
	}
	for _, w := range r.Warnings {
		if strings.Contains(w, "daily cap") {
			t.Fatalf("restricted capped: %v", r.Warnings)
		}
	}
	if db.GetSearchCredits(NameTavily, day) != 1 {
		t.Fatalf("restricted must not raise fleet row, got %d", db.GetSearchCredits(NameTavily, day))
	}
	if db.GetSearchCredits(NameTavily+":restricted", day) != 1 {
		t.Fatalf("restricted row %d", db.GetSearchCredits(NameTavily+":restricted", day))
	}
	fleet := e.Run(context.Background(), "q2", 8, false, "titan", dcache.ClassFleet)
	n := 0
	for _, w := range fleet.Warnings {
		if strings.Contains(w, "daily cap") {
			n++
		}
	}
	if n != 1 {
		t.Fatalf("fleet still capped: %v", fleet.Warnings)
	}
}

func TestDedupeLongerSnippet(t *testing.T) {
	in := []Hit{
		{URL: "https://Ex.com/a#x", Snippet: "aa", Source: NameSearxNG},
		{URL: "https://ex.com/a", Snippet: "aabbcc", Source: NameTavily},
	}
	out := Dedupe(in)
	if len(out) != 1 || out[0].Source != NameSearxNG || out[0].Snippet != "aabbcc" {
		t.Fatalf("%+v", out)
	}
	if out[0].URL != "https://Ex.com/a#x" {
		t.Fatalf("dedupe must keep first original URL, got %q", out[0].URL)
	}
}