package main

import (
	"context"
	"errors"
	"fmt"
	"maps"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"testing"
	"unicode/utf8"

	"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 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", longText(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)

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

	for level, count := range []int{12, 6, 3, 1} {
		for index := range count {
			require.True(t, memory.tree.Built(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", longText(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, memory.tree.Built(0, 8))
	summarizer.SetFailing(false)

	_, err = memory.Append("user", "retry")
	require.NoError(t, err)
	require.NoError(t, memory.Wait(t.Context()))
	require.True(t, memory.tree.Built(0, 8))
	require.True(t, memory.tree.Built(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", longText(0))
	require.NoError(t, err)
	require.NoError(t, memory.Wait(t.Context()))

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

	left, _ := memory.tree.Node(0, 0)
	right, _ := memory.tree.Node(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, memory.tree.Built(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", longText(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, memory.tree.Built(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, built := memory.tree.Node(2, 0)
	require.True(t, built)
	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", longText(0))
	require.NoError(t, err)
	require.NoError(t, memory.Wait(t.Context()))
	require.Equal(t, "<chat>\n0+1|"+pending+"\n</chat>", memory.View())

	summarizer.SetFailing(false)

	_, err = memory.Append("user", "again")
	require.NoError(t, err)
	require.NoError(t, memory.Wait(t.Context()))
	require.Equal(t, "<chat>\n0+1|"+fakeLine(memory.task(nodeKey{0, 0}))+"\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, errNoView)
}

// fakeSummarizer builds each node with a line of fakeLine. It fails each task
// while failing is true, and each task waits for gate when gate is not nil.
type fakeSummarizer struct {
	mutex   sync.Mutex
	gate    chan struct{}
	failing bool
	active  int
	peak    int
	tasks   []Task
}

func newFakeSummarizer() *fakeSummarizer {
	return &fakeSummarizer{}
}

func (f *fakeSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
	f.mutex.Lock()
	f.tasks = append(f.tasks, task)
	f.active++
	f.peak = max(f.peak, f.active)
	fail := f.failing
	f.mutex.Unlock()

	defer func() {
		f.mutex.Lock()
		f.active--
		f.mutex.Unlock()
	}()

	if f.gate != nil {
		select {
		case <-f.gate:
		case <-ctx.Done():
			return "", ctx.Err()
		}
	}

	if fail {
		return "", errors.New("fake failure")
	}

	return fakeLine(task), nil
}

// fakeLine returns a line of nodeBytes bytes or less that names the node and
// starts with the input. A real line of this size makes each merge a model
// call.
func fakeLine(task Task) string {
	count := 1 << task.Level
	line := fmt.Sprintf("summary %d+%d: %s", task.Index*count, count, strings.Join(task.Input, " "))
	line = strings.ReplaceAll(line, "\n", " ")
	for len(line) > nodeBytes-12 {
		_, size := utf8.DecodeLastRuneInString(line)
		line = line[:len(line)-size]
	}

	return line
}

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

// SetFailing makes the next tasks fail or not.
func (f *fakeSummarizer) SetFailing(failing bool) {
	f.mutex.Lock()
	defer f.mutex.Unlock()

	f.failing = failing
}

// Tasks returns the tasks so far.
func (f *fakeSummarizer) Tasks() []Task {
	f.mutex.Lock()
	defer f.mutex.Unlock()

	return slices.Clone(f.tasks)
}