package main

import (
	"context"
	"errors"
	"fmt"
	"maps"
	"math/rand/v2"
	"os"
	"path/filepath"
	"runtime"
	"slices"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

const (
	// replayMessages is the length of the synthetic log.
	replayMessages = 100000

	// replayShortMessages is the length of the synthetic log with -short.
	// It is long enough for batches of the chat view.
	replayShortMessages = 5000
)

func TestMemoryRestartLoadsTheSameView(t *testing.T) {
	root := t.TempDir()
	clock := newTestClock()
	memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	memory.chat.high = 200
	memory.chat.low = 100

	memory.cview.high = 80
	memory.cview.low = 40
	for n := range 100 {
		_, err = memory.Append("user", fmt.Sprint(n))
		require.NoError(t, err)
	}

	view := memory.View()
	context := memory.Context(0, 100)
	saved, err := os.ReadFile(filepath.Join(root, "view.json"))
	require.NoError(t, err)
	require.Less(t, len(memory.chat.lines), 100)
	require.NoError(t, memory.Close())

	memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	require.Equal(t, view, memory.View())
	require.Equal(t, context, memory.Context(0, 100))

	loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
	require.NoError(t, err)
	require.Equal(t, saved, loaded)
}

func TestMemoryOpensAfterACrashBeforeTheViewSave(t *testing.T) {
	root := t.TempDir()
	clock := newTestClock()
	memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	_, err = memory.Append("user", "first")
	require.NoError(t, err)

	saved, err := os.ReadFile(filepath.Join(root, "view.json"))
	require.NoError(t, err)

	_, err = memory.Append("user", "second")
	require.NoError(t, err)

	view := memory.View()
	require.NoError(t, memory.Close())

	// The log and the tree hold message 1, but the view does not.
	require.NoError(t, os.WriteFile(filepath.Join(root, "view.json"), saved, 0o600))

	memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)
	require.Equal(t, view, memory.View())
	require.NoError(t, memory.Close())
}

func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) {
	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	memory.chat.high = 200
	memory.chat.low = 100

	memory.cview.high = 80
	memory.cview.low = 40
	for n := range 64 {
		_, err = memory.Append("user", fmt.Sprint(n%10))
		require.NoError(t, err)

		coarse := 0
		for _, line := range memory.chat.lines {
			outer := memory.cview.lines[coarse]
			for nodeEnd(outer.level, outer.index) <= line.index<<line.level {
				coarse++
				outer = memory.cview.lines[coarse]
			}

			require.GreaterOrEqual(t, outer.level, line.level, "message %d", n)
		}
	}
}

func TestMemoryRunsAtMostEightBuilds(t *testing.T) {
	summarizer := newFakeSummarizer()
	summarizer.gate = make(chan struct{})
	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	for n := range 12 {
		_, err = memory.Append("echo", memoryLongText(n))
		require.NoError(t, err)
	}

	memory.mutex.Lock()

	running := slices.Collect(maps.Keys(memory.queue.running))
	memory.mutex.Unlock()
	slices.SortFunc(running, func(a, b NodeKey) int {
		return a.index - b.index
	})
	require.Equal(t, []NodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5}, {0, 6}, {0, 7}}, running)

	require.Eventually(t, func() bool {
		return summarizer.Active() == buildMax
	}, 10*time.Second, time.Millisecond)

	close(summarizer.gate)
	require.NoError(t, memory.Wait(t.Context()))
	require.Equal(t, buildMax, summarizer.peak)

	for level, count := range []int{12, 6, 3, 1} {
		for index := range count {
			require.True(t, treeIsBuilt(t, memory.tree, level, index), "node (%d, %d)", level, index)
		}
	}
}

func TestMemorySettleWaitsForThePendingLines(t *testing.T) {
	summarizer := newFakeSummarizer()
	summarizer.gate = make(chan struct{})
	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	_, err = memory.Append("echo", memoryLongText(0))
	require.NoError(t, err)
	require.Contains(t, memory.View(), viewPending)

	settled := make(chan error, 1)
	go func() {
		settled <- memory.Settle(t.Context())
	}()

	select {
	case err = <-settled:
		require.Fail(t, "Settle returned before the build", "error: %v", err)
	case <-time.After(50 * time.Millisecond):
	}

	close(summarizer.gate)
	require.NoError(t, <-settled)
	require.NotContains(t, memory.View(), viewPending)
}

func TestMemorySettleFailsWhileABuildFailed(t *testing.T) {
	summarizer := newFakeSummarizer()
	summarizer.SetFailing(true)

	logs := &SyncBuffer{}
	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, daemonTestLogger(logs))
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	_, err = memory.Append("echo", memoryLongText(0))
	require.NoError(t, err)
	require.ErrorIs(t, memory.Settle(t.Context()), errMemoryUnsettled)
	require.Equal(t, "<4>summary failed level=0 index=0: fake failure\n", logs.String())

	summarizer.SetFailing(false)

	_, err = memory.Append("user", "retry")
	require.NoError(t, err)
	require.NoError(t, memory.Settle(t.Context()))
}

func TestMemorySettleEndsWithItsContext(t *testing.T) {
	summarizer := newFakeSummarizer()
	summarizer.gate = make(chan struct{})
	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	_, err = memory.Append("echo", memoryLongText(0))
	require.NoError(t, err)

	ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
	defer cancel()

	require.ErrorIs(t, memory.Settle(ctx), context.DeadlineExceeded)
}

func TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) {
	summarizer := newFakeSummarizer()
	summarizer.SetFailing(true)

	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	for n := range 9 {
		_, err = memory.Append("echo", memoryLongText(n))
		require.NoError(t, err)
	}

	require.NoError(t, memory.Wait(t.Context()))

	for _, task := range summarizer.Tasks() {
		require.Less(t, task.Index, 8)
	}

	require.False(t, treeIsBuilt(t, memory.tree, 0, 8))
	summarizer.SetFailing(false)

	_, err = memory.Append("user", "retry")
	require.NoError(t, err)
	require.NoError(t, memory.Wait(t.Context()))
	require.True(t, treeIsBuilt(t, memory.tree, 0, 8))
	require.True(t, treeIsBuilt(t, memory.tree, 3, 0))
}

func TestMemoryMergesAfterBothHalves(t *testing.T) {
	summarizer := newFakeSummarizer()
	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	_, err = memory.Append("echo", memoryLongText(0))
	require.NoError(t, err)
	require.NoError(t, memory.Wait(t.Context()))

	_, err = memory.Append("echo", memoryLongText(1))
	require.NoError(t, err)
	require.NoError(t, memory.Wait(t.Context()))
	require.Len(t, summarizer.Tasks(), 3)

	left := treeReadNode(t, memory.tree, 0, 0)
	right := treeReadNode(t, memory.tree, 0, 1)
	merge := summarizer.Tasks()[2]
	require.Equal(t, []string{left.Text, right.Text}, merge.Input)
	require.Equal(t, "<chat>\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n</chat>", merge.Context)
	require.True(t, treeIsBuilt(t, memory.tree, 1, 0))
}

func TestMemoryRestartRestoresTheQueue(t *testing.T) {
	root := t.TempDir()
	clock := newTestClock()
	summarizer := newFakeSummarizer()
	summarizer.SetFailing(true)

	memory, err := openMemory(t.Context(), root, clock.Now, summarizer, discardLog())
	require.NoError(t, err)

	for n := range 4 {
		_, err = memory.Append("echo", memoryLongText(n))
		require.NoError(t, err)
	}

	require.NoError(t, memory.Wait(t.Context()))
	require.NoError(t, memory.Close())

	memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	require.NoError(t, memory.Wait(t.Context()))
	require.True(t, treeIsBuilt(t, memory.tree, 2, 0))
}

func TestMemoryBuildsFreeParents(t *testing.T) {
	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	for _, text := range []string{"a", "b", "c", "d"} {
		_, err = memory.Append("user", text)
		require.NoError(t, err)
	}

	node := treeReadNode(t, memory.tree, 2, 0)
	require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
}

func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {
	summarizer := newFakeSummarizer()
	summarizer.SetFailing(true)

	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	_, err = memory.Append("echo", memoryLongText(0))
	require.NoError(t, err)
	require.NoError(t, memory.Wait(t.Context()))
	require.Equal(t, "<chat>\n0+1|"+viewPending+"\n</chat>", memory.View())

	summarizer.SetFailing(false)

	_, err = memory.Append("user", "again")
	require.NoError(t, err)
	require.NoError(t, memory.Wait(t.Context()))

	task, err := memory.task(NodeKey{0, 0})
	require.NoError(t, err)
	require.Equal(t, "<chat>\n0+1|"+fakeLine(task)+"\n1+1|user: again\n</chat>", memory.View())
}

func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
	root := t.TempDir()
	clock := newTestClock()
	memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	_, err = memory.Append("user", "a")
	require.NoError(t, err)

	_, err = memory.log.Append("reply", "b")
	require.NoError(t, err)
	require.NoError(t, memory.Close())

	memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	require.Equal(t, "<chat>\n0+1|user: a\n1+1|reply: b\n</chat>", memory.View())
}

func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
	root := t.TempDir()
	clock := newTestClock()
	memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	_, err = memory.Append("user", "a")
	require.NoError(t, err)
	require.NoError(t, memory.Close())
	require.NoError(t, os.Remove(filepath.Join(root, "view.json")))

	_, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.ErrorIs(t, err, errMemoryNoView)
}

func TestQueuedAppendResumesWithoutADuplicate(t *testing.T) {
	dir := t.TempDir()
	memory, err := openMemory(t.Context(), dir, newTestClock().Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	var saved *int

	first, err := memory.AppendOnce("user", "Record the gate size.", nil, func(id int) error {
		saved = &id

		require.Equal(t, 0, memory.log.Len())

		return nil
	})
	require.NoError(t, err)
	require.NoError(t, memory.Close())

	memory, err = openMemory(t.Context(), dir, newTestClock().Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	second, err := memory.AppendOnce("user", first.Text, saved, func(int) error {
		require.FailNow(t, "a recorded message must not be prepared again")

		return nil
	})
	require.NoError(t, err)
	require.Equal(t, first, second)
	require.Equal(t, 1, memory.Len())
}

func TestQueuedAppendRetriesAnUnwrittenPreparedMessage(t *testing.T) {
	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	failure := errors.New("fake queue write failure")

	var saved *int

	_, err = memory.AppendOnce("work", "[Garden] Check the gate.", nil, func(id int) error {
		saved = &id

		return failure
	})
	require.ErrorIs(t, err, failure)
	require.Equal(t, 0, memory.Len())

	message, err := memory.AppendOnce("work", "[Garden] Check the gate.", saved, func(id int) error {
		require.Equal(t, *saved, id)

		return nil
	})
	require.NoError(t, err)
	require.Equal(t, 0, message.ID)
	require.Equal(t, 1, memory.Len())
}

func TestQueuedAppendRefusesAnInvalidReceipt(t *testing.T) {
	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	_, err = memory.Append("user", "Keep the original.")
	require.NoError(t, err)

	testCases := []struct {
		name  string
		id    int
		cause string
	}{
		{
			name:  "receipt points to different text",
			id:    0,
			cause: "does not match",
		},
		{
			name:  "receipt skips a log position",
			id:    2,
			cause: "beyond the log",
		},
		{
			name:  "receipt has a negative position",
			id:    -1,
			cause: "read queued message",
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			_, err := memory.AppendOnce("user", "Other text.", &tc.id, func(int) error {
				require.FailNow(t, "an invalid receipt must not be prepared")

				return nil
			})
			require.ErrorContains(t, err, tc.cause)
			require.Equal(t, 1, memory.Len())
		})
	}
}

// TestReplaySyntheticLog runs a synthetic log through the memory core with the
// fake summarizer. After each message it waits for the builds, as a turn does.
// It checks that each view stays within its sawtooth, that no view line holds a
// whole long message or a part of an unbuilt one, and that a restart in the
// middle and at the end loads the same views byte for byte. With -short, the
// log is shorter, and the test still requires a batch of the chat view.
func TestReplaySyntheticLog(t *testing.T) {
	start := time.Now()
	messages := replayMessages
	if testing.Short() {
		messages = replayShortMessages
	}

	root := t.TempDir()
	clock := newTestClock()
	random := rand.New(rand.NewPCG(1, 2))
	memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, memory.Close())
	}()

	batches := 0

	low := viewChatHigh
	long := make([]bool, messages)
	for id := range messages {
		if id == messages/2 {
			memory = replayRequireRestart(t, memory, root, clock)
		}

		clock.now = clock.now.Add(time.Minute)
		kind, text := replayMessage(random, id)
		long[id] = len(kind)+len(": ")+len(text) > nodeBytes
		lines := replayViewLines(memory)
		_, err = memory.Append(kind, text)
		require.NoError(t, err)

		size, merged := replayRequireViews(t, memory, lines, long, id%1000 == 0)
		if merged {
			batches++
			low = min(low, size)
		}

		require.NoError(t, memory.Wait(t.Context()))
	}

	memory = replayRequireRestart(t, memory, root, clock)
	require.Equal(t, messages, memory.log.Len())
	require.Positive(t, batches)

	require.NoError(t, memory.Wait(t.Context()))
	t.Logf("%d messages, %d nodes, %d batches in %s", messages, memory.tree.Count(), batches, time.Since(start).Round(time.Second))
	t.Logf("chat view: %d lines, %d bytes, lowest after a batch %d bytes", len(memory.chat.lines), memory.chat.Size(), low)
	t.Logf("compaction view: %d lines, %d bytes", len(memory.cview.lines), memory.cview.Size())

	runtime.GC()

	var stats runtime.MemStats

	runtime.ReadMemStats(&stats)
	t.Logf("heap in use: %d KB", stats.HeapInuse/1024)
}

// memoryLongText returns a text that is too long for its own node.
func memoryLongText(n int) string {
	return fmt.Sprintf("long %d ", n) + strings.Repeat("x", nodeBytes)
}

// replayRequireRestart closes memory, opens it again, and checks that the views
// and view.json stay the same byte for byte.
func replayRequireRestart(t *testing.T, memory *Memory, root string, clock *TestClock) *Memory {
	t.Helper()

	require.NoError(t, memory.Wait(t.Context()))

	view := memory.View()
	context := memory.Context(0, memory.log.Len())
	saved, err := os.ReadFile(filepath.Join(root, "view.json"))
	require.NoError(t, err)
	require.NoError(t, memory.Close())

	memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)
	require.NoError(t, memory.Wait(t.Context()))
	require.Equal(t, view, memory.View())
	require.Equal(t, context, memory.Context(0, memory.log.Len()))

	loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
	require.NoError(t, err)
	require.Equal(t, saved, loaded)

	return memory
}

// replayMessage returns the kind and the text of message id. Most messages are
// tool calls and their output. Users write short lines. One output in a hundred
// is over 30 KB.
func replayMessage(random *rand.Rand, id int) (string, string) {
	roll := random.IntN(100)
	switch {
	case roll < 5:
		return "user", replayText(id, 10+random.IntN(200))
	case roll < 10:
		return "reply", replayText(id, 100+random.IntN(1500))
	case roll < 12:
		return "note", replayText(id, 50+random.IntN(600))
	case roll < 13:
		return "work", replayText(id, 300+random.IntN(3000))
	case roll < 56:
		return "tool", replayText(id, 60+random.IntN(400))
	case roll < 57:
		return "echo", replayText(id, 30000+random.IntN(30000))
	default:
		return "echo", replayText(id, 20+random.IntN(4000))
	}
}

// replayViewLines returns the number of lines in the chat view of memory.
func replayViewLines(memory *Memory) int {
	memory.mutex.Lock()
	defer memory.mutex.Unlock()

	return len(memory.chat.lines)
}

// replayRequireViews checks the views after a new message, and with full
// compares them with the tree. Each view is at most its high limit, and a batch
// merges the chat view to its low limit. The leaf of the new message can be
// built after the append, so the sizes count it as the pending mark, as the
// append did. It returns the size of the chat view and if a batch merged it,
// from the number of lines before the message.
func replayRequireViews(t *testing.T, memory *Memory, lines int, long []bool, full bool) (int, bool) {
	t.Helper()

	memory.mutex.Lock()
	defer memory.mutex.Unlock()

	size := replayAppendSize(memory.chat)
	require.LessOrEqual(t, size, viewChatHigh)
	require.LessOrEqual(t, replayAppendSize(memory.cview), viewCompactionHigh)

	merged := len(memory.chat.lines) < lines
	if merged {
		require.LessOrEqual(t, size, viewChatLow)
	}

	replayRequireOnlySummaries(t, memory, long, full)

	return size, merged
}

// replayText returns a text of about size bytes that starts with its message ID
// and holds some non-ASCII characters.
func replayText(id, size int) string {
	head := fmt.Sprintf("m%d ", id)
	words := "ação data line café "

	return head + strings.Repeat(words, max(1, (size-len(head))/len(words)))
}

// replayAppendSize returns the size of view with its last line as the pending
// mark.
func replayAppendSize(view *View) int {
	last := view.lines[len(view.lines)-1]
	return view.Size() - view.lineSize(last) + len(viewPending)
}

// replayRequireOnlySummaries checks each line of the chat view as the model
// reads it. A built line holds its node, a long message has a summary of the
// fake summarizer and not its text, and an unbuilt line holds only the pending
// mark. With full, it also checks that the view keeps the node of each built
// line. The caller holds the lock of memory.
func replayRequireOnlySummaries(t *testing.T, memory *Memory, long []bool, full bool) {
	t.Helper()

	rendered := strings.Split(memory.chat.Render(), "\n")
	require.Len(t, rendered, len(memory.chat.lines)+2)

	for n, line := range memory.chat.lines {
		head := fmt.Sprintf("%d+%d|", line.index<<line.level, 1<<line.level)
		node, built := memory.chat.nodes[line]
		if full {
			require.Equal(t, treeIsBuilt(t, memory.tree, line.level, line.index), built)
		}

		if !built {
			require.Equal(t, 0, line.level)
			require.Equal(t, head+viewPending, rendered[n+1])

			continue
		}

		require.Equal(t, head+strings.ReplaceAll(node.Text, "\n", " "), rendered[n+1])

		if line.level == 0 && long[line.index] {
			require.True(t, strings.HasPrefix(node.Text, "summary "), "message %d", line.index)
		}
	}
}