package digest

// P5b owner review: regression tests for the background embedder.

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"
)

// flakyDixie is fakeDixie whose /embed answers 500 to its first failN requests.
func flakyDixie(t *testing.T, failN int64) (*fakeDixie, *atomic.Int64) {
	f := &fakeDixie{}
	var calls atomic.Int64
	var disc atomic.Bool
	mux := http.NewServeMux()
	mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) {
		if calls.Add(1) <= failN {
			http.Error(w, "boom", http.StatusInternalServerError)
			return
		}
		var body struct {
			Input []string `json:"input"`
		}
		_ = json.NewDecoder(r.Body).Decode(&body)
		var data []map[string]any
		for i, s := range body.Input {
			data = append(data, map[string]any{"index": i, "embedding": bow(s)})
		}
		f.embedded.Add(int64(len(body.Input)))
		_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
	})
	mux.Handle("/chat", sseHandler(t, "ok", &f.lastChat, &disc))
	f.srv = httptest.NewServer(mux)
	t.Cleanup(f.srv.Close)
	return f, &calls
}

// A document whose background embedding failed is queued again by a later question; it must
// not stay "seen" (and so BM25-only) for the life of the process.
func TestBackgroundRequeuesAfterFailure(t *testing.T) {
	f, _ := flakyDixie(t, 1)
	e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
	doc := longDoc(300, 250)
	n := len(SplitPages(doc, ChunkTokens))
	wait := func() {
		t.Helper()
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		if err := e.WaitBackground(ctx); err != nil {
			t.Fatal(err)
		}
	}
	q := "who signed the zanzibarite manifest?"
	if r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet"); r.Retrieval != "bm25" {
		t.Fatalf("first: %+v", r)
	}
	wait() // the first batch failed: the document was dropped
	if r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet"); r.Retrieval != "bm25" {
		t.Fatalf("second: %+v, want bm25 (nothing embedded yet)", r)
	}
	wait() // re-queued by the second question, and embedded this time
	r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet")
	if r.Retrieval != "hybrid" || r.Embedded != n {
		t.Fatalf("third: %+v, want hybrid with all %d embedded", r, n)
	}
}

// The background worker waits for a busy gate; it does not give up on the document.
func TestBackgroundWaitsForBusyGate(t *testing.T) {
	f := newFakeDixie(t, false)
	gate := NewGate(1, 0, 50*time.Millisecond) // no queue: a busy gate refuses at once
	e := f.engine(t, gate, true)
	doc := longDoc(100, 0)
	n := len(SplitPages(doc, ChunkTokens))
	if ok, _ := gate.Acquire(context.Background()); !ok {
		t.Fatal("could not take the gate")
	}
	// Run takes the gate before it chooses a retrieval mode, so queue the document directly.
	chunks := SplitPages(doc, ChunkTokens)
	e.queueForEmbed(chunks, "hash-g", "fleet")
	time.Sleep(2500 * time.Millisecond) // longer than the old single retry
	gate.Release()
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	if err := e.WaitBackground(ctx); err != nil {
		t.Fatal(err)
	}
	if got := f.embedded.Load(); got != int64(n) {
		t.Fatalf("embedded %d of %d chunks: the worker gave up on a busy gate", got, n)
	}
}