package memory

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"

	"github.com/shaikn6/lethe/internal/policy"
)

var errDisk = errors.New("disk error")

// flakyFile fails writes on demand, writing half the bytes first the way a
// full or failing disk does.
type flakyFile struct {
	journalFile
	failWrites   bool
	failTruncate bool
}

func (f *flakyFile) Write(p []byte) (int, error) {
	if f.failWrites {
		n, _ := f.journalFile.Write(p[:len(p)/2])
		return n, errDisk
	}
	return f.journalFile.Write(p)
}

func (f *flakyFile) Truncate(size int64) error {
	if f.failTruncate {
		return errDisk
	}
	return f.journalFile.Truncate(size)
}

func breakDisk(s *Store) *flakyFile {
	flaky := &flakyFile{journalFile: s.journal.file, failWrites: true}
	s.journal.file = flaky
	return flaky
}

func TestFailedWriteIsRolledBackAndTheJournalStaysReadable(t *testing.T) {
	path := filepath.Join(t.TempDir(), "store.journal")
	s := open(t, Options{JournalPath: path})
	remember(t, s, Memory{Text: "written before the disk problem"})

	disk := breakDisk(s)
	if _, err := s.Remember(Memory{Namespace: "ns", Text: "lost to the disk problem"}); !errors.Is(err, errDisk) {
		t.Fatalf("Remember on a failing disk = %v, want the disk error", err)
	}
	if got := recall(t, s, Query{Text: "lost disk problem"}); len(got.Memories) != 1 {
		t.Fatalf("a memory that was never written is being served: %v", texts(got))
	}
	disk.failWrites = false
	remember(t, s, Memory{Text: "written after the disk recovered"})
	if st := s.Stats(); st.Memories != 2 || st.JournalErrors != 1 || !st.JournalHealthy {
		t.Fatalf("stats = %+v", st)
	}
	s.Close()

	reopened := open(t, Options{JournalPath: path})
	if st := reopened.Stats(); st.Memories != 2 {
		t.Fatalf("after restart: %d memories, want 2 (journal must not hold a half-written line)", st.Memories)
	}
}

func TestJournalThatCannotBeRolledBackRefusesFurtherWrites(t *testing.T) {
	s := open(t, Options{JournalPath: filepath.Join(t.TempDir(), "store.journal")})
	kept := remember(t, s, Memory{Text: "safe memory"})
	disk := breakDisk(s)
	disk.failTruncate = true
	if _, err := s.Remember(Memory{Namespace: "ns", Text: "x"}); err == nil {
		t.Fatal("Remember succeeded on a failing disk")
	}
	disk.failWrites, disk.failTruncate = false, false
	if _, err := s.Remember(Memory{Namespace: "ns", Text: "y"}); err == nil || !strings.Contains(err.Error(), "unusable") {
		t.Fatalf("Remember after an unrecoverable journal error = %v, want it refused", err)
	}
	if err := s.Forget("ns", kept.ID); err == nil {
		t.Fatal("Forget succeeded with an unusable journal")
	}
	if st := s.Stats(); st.JournalHealthy || st.Memories != 1 {
		t.Fatalf("stats = %+v, want unhealthy journal and the one safe memory", st)
	}
	if got := recall(t, s, Query{Text: "safe memory"}); len(got.Memories) != 1 {
		t.Fatalf("reads must keep working: %v", texts(got))
	}
}

func TestFeedbackAndWeightsNeverRunAheadOfTheJournal(t *testing.T) {
	s := open(t, Options{JournalPath: filepath.Join(t.TempDir(), "store.journal")})
	for i := range 5 {
		remember(t, s, Memory{Text: fmt.Sprintf("note %d about the topic", i), Embedding: []float32{float32(i), 1}})
	}
	r := recall(t, s, Query{Text: "topic note", Embedding: []float32{1, 1}, Explore: true})
	before := s.Stats()

	disk := breakDisk(s)
	if err := s.Feedback(r.ID, 1, nil); !errors.Is(err, errDisk) {
		t.Fatalf("Feedback on a failing disk = %v, want the disk error", err)
	}
	var other policy.Features
	other[policy.Recency] = 3
	if err := s.SetWeights(other); !errors.Is(err, errDisk) {
		t.Fatalf("SetWeights on a failing disk = %v, want the disk error", err)
	}
	after := s.Stats()
	if after.Weights != before.Weights || after.Updates != before.Updates || after.Baseline != before.Baseline || after.Pending != 1 {
		t.Fatalf("state changed although nothing was journaled:\n got %+v\nwant %+v", after, before)
	}

	disk.failWrites = false
	if err := s.Feedback(r.ID, 1, nil); err != nil {
		t.Fatalf("retry after the disk recovered = %v", err)
	}
	if st := s.Stats(); st.Updates != before.Updates+1 || st.Pending != 0 {
		t.Fatalf("stats after retry = %+v", st)
	}
}

func TestForgottenTextIsErasedFromMemoryAndDisk(t *testing.T) {
	path := filepath.Join(t.TempDir(), "store.journal")
	s := open(t, Options{JournalPath: path})
	remember(t, s, Memory{Text: "the gate code is on the notice board"})
	secret := remember(t, s, Memory{Text: "card number zebra4417 expires soon", Metadata: map[string]string{"pin": "quokka9921"}})
	onDisk := func() string {
		data, err := os.ReadFile(path)
		if err != nil {
			t.Fatal(err)
		}
		return string(data)
	}
	if !strings.Contains(onDisk(), "zebra4417") {
		t.Fatal("test setup: the secret should be in the journal before it is forgotten")
	}
	if err := s.Forget("ns", secret.ID); err != nil {
		t.Fatal(err)
	}
	sp := s.spaces["ns"]
	for _, e := range sp.entries {
		if strings.Contains(e.Text, "zebra4417") || e.Metadata["pin"] != "" {
			t.Fatalf("forgotten memory is still in RAM: %+v", e.Memory)
		}
	}
	if got := sp.ix.Search("zebra4417", 5); len(got) != 0 {
		t.Fatal("forgotten word is still searchable")
	}
	if err := s.Compact(); err != nil {
		t.Fatal(err)
	}
	for _, leaked := range []string{"zebra4417", "quokka9921", secret.ID} {
		if strings.Contains(onDisk(), leaked) {
			t.Fatalf("%q is still in the journal after compaction", leaked)
		}
	}
	if !strings.Contains(onDisk(), "gate code") {
		t.Fatal("compaction dropped a live memory")
	}
}

func TestForgottenTextDoesNotSurviveARestartOrManyForgets(t *testing.T) {
	path := filepath.Join(t.TempDir(), "store.journal")
	s := open(t, Options{JournalPath: path, CompactAfter: 3})
	var ids []string
	for i := range 3 {
		ids = append(ids, remember(t, s, Memory{Text: fmt.Sprintf("disposable wombat%d note", i)}).ID)
	}
	keep := remember(t, s, Memory{Text: "permanent note"})
	for _, id := range ids {
		if err := s.Forget("ns", id); err != nil {
			t.Fatal(err)
		}
	}
	data, _ := os.ReadFile(path)
	if bytes.Contains(data, []byte("wombat")) {
		t.Fatal("journal still holds forgotten text after CompactAfter forgets")
	}

	late := remember(t, s, Memory{Text: "another disposable numbat note"})
	if err := s.Forget("ns", late.ID); err != nil {
		t.Fatal(err)
	}
	s.Close()
	data, _ = os.ReadFile(path)
	if bytes.Contains(data, []byte("numbat")) {
		t.Fatal("journal still holds forgotten text after Close")
	}
	reopened := open(t, Options{JournalPath: path})
	got := recall(t, reopened, Query{Text: "permanent note"})
	if len(got.Memories) != 1 || got.Memories[0].ID != keep.ID {
		t.Fatalf("live memory lost across compaction and restart: %v", texts(got))
	}
}

func TestCompactionKeepsRankingAndLearningAndShrinksTheJournal(t *testing.T) {
	path := filepath.Join(t.TempDir(), "store.journal")
	s := open(t, Options{JournalPath: path, CompactSlack: 1})
	for i := range 40 {
		remember(t, s, Memory{Text: fmt.Sprintf("meeting note %d about the budget", i), Embedding: []float32{float32(i % 7), float32(i % 5), 1}, CreatedAt: day0.AddDate(0, 0, i)})
	}
	for i := range 300 {
		r := recall(t, s, Query{Text: "latest budget meeting note", Embedding: []float32{1, 2, 1}, Explore: true})
		if err := s.Feedback(r.ID, float64(i%2), nil); err != nil {
			t.Fatal(err)
		}
	}
	query := Query{Text: "latest budget meeting note", Embedding: []float32{1, 2, 1}, K: 10}
	before, stats := recall(t, s, query), s.Stats()
	sizeBefore := fileSize(t, path)
	if err := s.Compact(); err != nil {
		t.Fatal(err)
	}
	if sizeAfter := fileSize(t, path); sizeAfter*3 > sizeBefore {
		t.Fatalf("journal went from %d to %d bytes; expected it to shrink by well over half", sizeBefore, sizeAfter)
	}
	sameRecall(t, "after compaction", recall(t, s, query), before)
	remember(t, s, Memory{Text: "written after the compaction"})
	s.Close()

	reopened := open(t, Options{JournalPath: path})
	sameRecall(t, "after compaction and restart", recall(t, reopened, query), before)
	if got := reopened.Stats(); got.Memories != 41 || got.Weights != stats.Weights || got.Updates != stats.Updates {
		t.Fatalf("stats after restart = %+v, want 41 memories and the learned state %+v", got, stats)
	}
}

func TestStoreRefusesNewMemoriesWhenFull(t *testing.T) {
	s := open(t, Options{MaxMemories: 2})
	first := remember(t, s, Memory{Text: "one"})
	remember(t, s, Memory{Namespace: "other", Text: "two"})
	if _, err := s.Remember(Memory{Namespace: "ns", Text: "three"}); !errors.Is(err, ErrFull) {
		t.Fatalf("Remember beyond the limit = %v, want ErrFull", err)
	}
	if err := s.Forget("ns", first.ID); err != nil {
		t.Fatal(err)
	}
	remember(t, s, Memory{Text: "three fits after one was forgotten"})
}

func TestRejectsOversizedMetadata(t *testing.T) {
	s := open(t, Options{})
	for name, metadata := range map[string]map[string]string{
		"long key":   {strings.Repeat("k", MaxMetadataKeyLen+1): "v"},
		"long value": {"k": strings.Repeat("v", MaxMetadataValLen+1)},
	} {
		if _, err := s.Remember(Memory{Namespace: "ns", Text: "x", Metadata: metadata}); !errors.Is(err, ErrInvalid) {
			t.Errorf("%s: %v, want ErrInvalid", name, err)
		}
	}
}

func TestCompactionDuringConcurrentUse(t *testing.T) {
	path := filepath.Join(t.TempDir(), "store.journal")
	s := open(t, Options{JournalPath: path, CompactAfter: 20})
	const workers, rounds = 8, 120
	var wg sync.WaitGroup
	var compactWG sync.WaitGroup
	var mu sync.Mutex
	kept := 0
	stop := make(chan struct{})
	compactWG.Add(1)
	go func() {
		defer compactWG.Done()
		for {
			select {
			case <-stop:
				return
			default:
				if err := s.Compact(); err != nil {
					t.Error(err)
				}
			}
		}
	}()
	for w := range workers {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ns := fmt.Sprintf("ns%d", w%3)
			for i := range rounds {
				m, err := s.Remember(Memory{Namespace: ns, Text: fmt.Sprintf("worker %d item %d on the shared subject", w, i), Embedding: []float32{float32(w), float32(i), 1}})
				if err != nil {
					t.Error(err)
					return
				}
				r, err := s.Recall(Query{Namespace: ns, Text: "shared subject item", Embedding: []float32{1, 1, 1}, K: 5, Explore: true})
				if err != nil {
					t.Error(err)
					return
				}
				for _, got := range r.Memories {
					if got.Text == "" {
						t.Error("recall returned an erased memory")
					}
				}
				if err := s.Feedback(r.ID, float64(i%2), nil); err != nil {
					t.Error(err)
					return
				}
				if i%3 == 0 {
					if err := s.Forget(ns, m.ID); err != nil {
						t.Error(err)
						return
					}
				} else {
					mu.Lock()
					kept++
					mu.Unlock()
				}
			}
		}()
	}
	wg.Wait()
	close(stop)
	compactWG.Wait() // join the compactor before Close so it cannot call Compact on a closed store
	st := s.Stats()
	if st.Memories != kept || st.Updates != workers*rounds || st.JournalErrors != 0 {
		t.Fatalf("stats = %+v, want %d memories and %d updates", st, kept, workers*rounds)
	}
	s.Close()
	reopened := open(t, Options{JournalPath: path})
	if got := reopened.Stats(); got.Memories != kept || got.Updates != st.Updates || got.Weights != st.Weights {
		t.Fatalf("after restart: %+v\nwant %d memories, state %+v", got, kept, st)
	}
}

func fileSize(t *testing.T, path string) int64 {
	t.Helper()
	info, err := os.Stat(path)
	if err != nil {
		t.Fatal(err)
	}
	return info.Size()
}

func sameRecall(t *testing.T, when string, got, want Recall) {
	t.Helper()
	if len(got.Memories) != len(want.Memories) {
		t.Fatalf("%s: %d memories, want %d", when, len(got.Memories), len(want.Memories))
	}
	for i := range got.Memories {
		if got.Memories[i].ID != want.Memories[i].ID || got.Memories[i].Score != want.Memories[i].Score {
			t.Fatalf("%s: position %d is %+v, want %+v", when, i, got.Memories[i], want.Memories[i])
		}
	}
}

func TestOpenRemovesAnInterruptedCompactionAndReportsUnerasedForgets(t *testing.T) {
	path := filepath.Join(t.TempDir(), "store.journal")
	s := open(t, Options{JournalPath: path})
	gone := remember(t, s, Memory{Text: "pangolin55 must not survive"})
	remember(t, s, Memory{Text: "this one stays"})
	if err := s.Forget("ns", gone.ID); err != nil {
		t.Fatal(err)
	}
	if st := s.Stats(); st.ForgottenOnDisk != 1 {
		t.Fatalf("ForgottenOnDisk = %d, want 1 before compaction", st.ForgottenOnDisk)
	}
	// Simulate a crash: no Close, and a half-finished compaction copy on disk.
	if err := os.WriteFile(path+compactSuffix, []byte(`{"remember":{"text":"pangolin55 stale copy"}}`), 0o600); err != nil {
		t.Fatal(err)
	}
	s.journal.close()

	reopened := open(t, Options{JournalPath: path})
	if _, err := os.Stat(path + compactSuffix); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("interrupted compaction copy still exists: %v", err)
	}
	data, _ := os.ReadFile(path)
	if bytes.Contains(data, []byte("pangolin55")) {
		t.Fatal("forgotten text survived a restart")
	}
	if st := reopened.Stats(); st.ForgottenOnDisk != 0 || st.Memories != 1 {
		t.Fatalf("stats after restart = %+v", st)
	}
}

func TestCompactIsANoOpWhenThereIsNothingToDo(t *testing.T) {
	path := filepath.Join(t.TempDir(), "store.journal")
	s := open(t, Options{JournalPath: path})
	remember(t, s, Memory{Text: "steady state"})
	before, err := os.Stat(path)
	if err != nil {
		t.Fatal(err)
	}
	for range 50 {
		if err := s.Compact(); err != nil {
			t.Fatal(err)
		}
	}
	after, _ := os.Stat(path)
	if !os.SameFile(before, after) {
		t.Fatal("Compact rewrote the journal although nothing had been forgotten")
	}
}

func TestCompactionDoesNotWriteThroughAPlantedLink(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "store.journal")
	victim := filepath.Join(dir, "victim.txt")
	if err := os.WriteFile(victim, []byte("do not overwrite"), 0o600); err != nil {
		t.Fatal(err)
	}
	s := open(t, Options{JournalPath: path})
	m := remember(t, s, Memory{Text: "to be forgotten"})
	remember(t, s, Memory{Text: "to be kept"})
	if err := os.Symlink(victim, path+compactSuffix); err != nil {
		t.Fatal(err)
	}
	if err := s.Forget("ns", m.ID); err != nil {
		t.Fatal(err)
	}
	if err := s.Compact(); err != nil {
		t.Fatal(err)
	}
	if data, _ := os.ReadFile(victim); string(data) != "do not overwrite" {
		t.Fatalf("compaction wrote through the link: victim now holds %q", data)
	}
	if info, err := os.Lstat(path); err != nil || info.Mode()&os.ModeSymlink != 0 {
		t.Fatalf("journal is not a regular file after compaction: %v %v", info, err)
	}
}