package main

import (
	"fmt"
	"maps"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"
	"time"

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

func TestMemoryRestartLoadsTheSameView(t *testing.T) {
	root := t.TempDir()
	clock := newTestClock()
	memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
	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())
	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())
	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())
	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())
	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)
	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 TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) {
	summarizer := newFakeSummarizer()
	summarizer.SetFailing(true)
	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
	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)
	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)
	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())
	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())
	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)
	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())
	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())
	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())
	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())
	require.ErrorIs(t, err, errMemoryNoView)
}

// 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)
}