package main

import (
	"fmt"
	"math/rand/v2"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

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

// syntheticMessages is the length of the synthetic log.
const syntheticMessages = 100000

// 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.
func TestReplaySyntheticLog(t *testing.T) {
	start := time.Now()
	root := t.TempDir()
	clock := newTestClock()
	random := rand.New(rand.NewPCG(1, 2))
	memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
	require.NoError(t, err)

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

	batches := 0
	low := chatHigh
	long := make([]bool, syntheticMessages)
	for id := range syntheticMessages {
		if id == syntheticMessages/2 {
			memory = requireRestart(t, memory, root, clock)
		}

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

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

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

	memory = requireRestart(t, memory, root, clock)
	require.Equal(t, syntheticMessages, memory.log.Len())
	require.Greater(t, batches, 0)

	require.NoError(t, memory.Wait(t.Context()))
	t.Logf("%d messages, %d nodes, %d batches in %s", syntheticMessages, 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)
}

// syntheticMessage 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 syntheticMessage(random *rand.Rand, id int) (string, string) {
	roll := random.IntN(100)
	switch {
	case roll < 5:
		return "user", syntheticText(id, 10+random.IntN(200))
	case roll < 10:
		return "reply", syntheticText(id, 100+random.IntN(1500))
	case roll < 12:
		return "note", syntheticText(id, 50+random.IntN(600))
	case roll < 13:
		return "work", syntheticText(id, 300+random.IntN(3000))
	case roll < 56:
		return "tool", syntheticText(id, 60+random.IntN(400))
	case roll < 57:
		return "echo", syntheticText(id, 30000+random.IntN(30000))
	default:
		return "echo", syntheticText(id, 20+random.IntN(4000))
	}
}

// syntheticText returns a text of about size bytes that starts with its message
// ID and holds some non-ASCII characters.
func syntheticText(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)))
}

// requireViews 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 requireViews(t *testing.T, memory *Memory, lines int, long []bool, full bool) (int, bool) {
	t.Helper()

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

	size := appendSize(memory.chat)
	require.LessOrEqual(t, size, chatHigh)
	require.LessOrEqual(t, appendSize(memory.cview), compactionHigh)

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

	requireOnlySummaries(t, memory, long, full)

	return size, merged
}

// requireOnlySummaries 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 requireOnlySummaries(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, isBuilt(t, memory.tree, line.level, line.index), built)
		}

		if !built {
			require.Equal(t, 0, line.level)
			require.Equal(t, head+pending, 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)
		}
	}
}

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

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

	return len(memory.chat.lines)
}

// requireRestart closes memory, opens it again, and checks that the views and
// view.json stay the same byte for byte.
func requireRestart(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())
	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
}