// Copyright (c) 2026 Sidero Labs, Inc.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.

package sqlitelog_test

import (
	"context"
	"errors"
	"fmt"
	"io"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/cosi-project/runtime/pkg/state"
	"github.com/cosi-project/runtime/pkg/state/impl/inmem"
	"github.com/cosi-project/runtime/pkg/state/impl/namespaced"
	"github.com/cosi-project/state-sqlite/pkg/sqlitexx"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/goleak"
	"go.uber.org/zap"
	"go.uber.org/zap/zaptest"
	"golang.org/x/sync/errgroup"

	"github.com/siderolabs/omni/client/pkg/omni/resources/omni"
	"github.com/siderolabs/omni/internal/backend/runtime/omni/sqlite"
	"github.com/siderolabs/omni/internal/pkg/config"
	"github.com/siderolabs/omni/internal/pkg/siderolink/logstore"
	"github.com/siderolabs/omni/internal/pkg/siderolink/logstore/sqlitelog"
)

func TestMain(m *testing.M) {
	goleak.VerifyTestMain(m)
}

func TestReadWrite(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)

	// Disable cleanup to avoid connection pool contention between
	// writers and cleanup queries. Cleanup correctness is covered by TestCleanupProbabilities.
	storageConf := config.Default().Logs.Machine.Storage
	storageConf.SetCleanupProbability(0)

	storeManager, _ := setupDBWithStorageConfig(ctx, t, logger, storageConf)

	sqliteStore1, err := storeManager.Create("test-1")
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, sqliteStore1.Close())
	})

	sqliteStore2, err := storeManager.Create("test-2")
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, sqliteStore2.Close())
	})

	numLines := 1000

	for i := range numLines {
		require.NoError(t, sqliteStore1.WriteLine(ctx, fmt.Appendf(nil, "Hello, World %d!", i)))
		require.NoError(t, sqliteStore2.WriteLine(ctx, fmt.Appendf(nil, "Hello, World %d!", i)))
	}

	t.Run("read all", func(t *testing.T) {
		t.Parallel()

		testRead(ctx, t, sqliteStore1, numLines, -1)
	})

	t.Run("tail", func(t *testing.T) {
		t.Parallel()

		testRead(ctx, t, sqliteStore1, 100, 100)
	})
}

func TestFollow(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)
	storeManager, _ := setupDB(ctx, t, logger)

	store, err := storeManager.Create("test-1")
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, store.Close())
	})

	require.NoError(t, store.WriteLine(ctx, []byte("Hello, World 1!")))
	require.NoError(t, store.WriteLine(ctx, []byte("Hello, World 2!")))

	rdr, err := store.Reader(ctx, -1, true)
	require.NoError(t, err)

	var wg sync.WaitGroup

	lineCh := make(chan string)

	wg.Go(func() {
		readLines(ctx, t, rdr, lineCh)
	})

	assertLine(ctx, t, lineCh, "Hello, World 1!")
	assertLine(ctx, t, lineCh, "Hello, World 2!")

	require.NoError(t, store.WriteLine(ctx, []byte("Hello, World 3!")))
	require.NoError(t, store.WriteLine(ctx, []byte("Hello, World 4!")))

	assertLine(ctx, t, lineCh, "Hello, World 3!")
	assertLine(ctx, t, lineCh, "Hello, World 4!")

	cancel()

	wg.Wait()
}

func TestTruncation(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)
	storeManager, _ := setupDB(ctx, t, logger)

	// 1. Setup Long ID and Huge Message
	longID := "machine-" + string(make([]byte, 200))

	hugeMessage := make([]byte, 20*1024)
	for i := range hugeMessage {
		hugeMessage[i] = 'a'
	}

	store, err := storeManager.Create(longID)
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, store.Close())
	})

	// 2. Write the huge message immediately.
	// This populates the DB so Exists() can return true,
	// and allows us to test message truncation simultaneously.
	err = store.WriteLine(ctx, hugeMessage)
	require.NoError(t, err)

	// 3. Test Machine ID Truncation (in Manager)
	// The manager should find the store even if we pass the original long ID.
	exists, err := storeManager.Exists(ctx, longID)
	require.NoError(t, err)
	assert.True(t, exists, "Manager should find store with truncated ID (requires at least one log line written)")

	// 4. Test Message Truncation (in Store)
	// Read it back.
	rdr, err := store.Reader(ctx, -1, false)
	require.NoError(t, err)

	lines := readAllLines(ctx, t, rdr)
	require.Len(t, lines, 1)

	// It should be exactly 16KB (16 * 1024).
	assert.Len(t, lines[0], 16*1024, "Message should be truncated to 16KB")
}

func TestManagerLifecycle(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)
	storeManager, _ := setupDB(ctx, t, logger)
	id := "lifecycle-test"

	// 1. Check Exists before creation.
	exists, err := storeManager.Exists(ctx, id)
	require.NoError(t, err)
	assert.False(t, exists)

	// 2. Create and write.
	store, err := storeManager.Create(id)
	require.NoError(t, err)
	require.NoError(t, store.WriteLine(ctx, []byte("log data")))

	require.NoError(t, store.Close())

	// 3. Check Exists after creation and close.
	exists, err = storeManager.Exists(ctx, id)
	require.NoError(t, err)
	assert.True(t, exists)

	// 4. Remove.
	err = storeManager.Remove(ctx, id)
	require.NoError(t, err)

	// 5. Check Exists after removal.
	exists, err = storeManager.Exists(ctx, id)
	require.NoError(t, err)
	assert.False(t, exists)

	// 6. Verify data is actually gone.
	storeRecreated, err := storeManager.Create(id)
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, storeRecreated.Close())
	})

	rdr, err := storeRecreated.Reader(ctx, -1, false)
	require.NoError(t, err)
	lines := readAllLines(ctx, t, rdr)
	assert.Empty(t, lines)
}

func TestReaderParameters(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)
	storeManager, _ := setupDB(ctx, t, logger)
	store, err := storeManager.Create("reader-params")
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, store.Close())
	})

	// Write 10 lines.
	for i := range 10 {
		require.NoError(t, store.WriteLine(ctx, fmt.Appendf(nil, "line %d", i)))
	}

	tests := []struct {
		name      string
		firstLine string
		nLines    int
		wantCount int
	}{
		{
			name:      "Read All (Negative)",
			nLines:    -1,
			wantCount: 10,
			firstLine: "line 0",
		},
		{
			name:      "Read None (Zero)",
			nLines:    0,
			wantCount: 0,
		},
		{
			name:      "Read Tail (Small)",
			nLines:    3,
			wantCount: 3,
			firstLine: "line 7",
		},
		{
			name:      "Read Tail (Large)",
			nLines:    100,
			wantCount: 10,
			firstLine: "line 0",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			rdr, err := store.Reader(ctx, tt.nLines, false)
			require.NoError(t, err)

			lines := readAllLines(ctx, t, rdr)
			assert.Len(t, lines, tt.wantCount)

			if tt.wantCount > 0 && tt.firstLine != "" {
				assert.Equal(t, tt.firstLine, lines[0])
			}
		})
	}
}

func TestFollowNoHistory(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)
	storeManager, _ := setupDB(ctx, t, logger)
	store, err := storeManager.Create("follow-no-history")
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, store.Close())
	})

	// 1. Write historical data.
	require.NoError(t, store.WriteLine(ctx, []byte("history 1")))
	require.NoError(t, store.WriteLine(ctx, []byte("history 2")))

	// 2. Request 0 lines, but follow=true.
	rdr, err := store.Reader(ctx, 0, true)
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, rdr.Close())
	})

	// 3. Write new data.
	expected := "new 1"
	require.NoError(t, store.WriteLine(ctx, []byte(expected)))

	// 4. Read the first line.
	// Logic check: nLines=0 should skip "history 1/2" and stream "new 1" immediately.
	line, err := rdr.ReadLine(ctx)
	require.NoError(t, err)

	assert.Equal(t, expected, string(line), "Reader(0, true) should ignore history and only return new logs")
}

func TestFollowTail(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)
	storeManager, _ := setupDB(ctx, t, logger)
	store, err := storeManager.Create("follow-tail")
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, store.Close())
	})

	// 1. Write 20 lines of history
	for i := range 20 {
		require.NoError(t, store.WriteLine(ctx, fmt.Appendf(nil, "history %d", i)))
	}

	// 2. Request last 5 lines and follow
	rdr, err := store.Reader(ctx, 5, true)
	require.NoError(t, err)

	lineCh := make(chan string)

	var wg sync.WaitGroup

	wg.Go(func() {
		readLines(ctx, t, rdr, lineCh)
	})

	// 3. Verify we get the last 5 lines of history (15-19)
	for i := 15; i < 20; i++ {
		assertLine(ctx, t, lineCh, fmt.Sprintf("history %d", i))
	}

	// 4. Write new data
	require.NoError(t, store.WriteLine(ctx, []byte("new 1")))

	// 5. Verify we get the new line
	assertLine(ctx, t, lineCh, "new 1")

	cancel()
	wg.Wait()
}

// TestFollowRapidWrites tests that a reader following rapid writes
// receives all lines without gaps or duplication.
func TestFollowRapidWrites(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)

	// Use a config with cleanup disabled to avoid connection pool contention
	// between the writer, reader, and cleanup queries. Cleanup correctness
	// is covered by TestCleanupProbabilities.
	storageConf := config.Default().Logs.Machine.Storage
	storageConf.SetCleanupProbability(0)

	storeManager, _ := setupDBWithStorageConfig(ctx, t, logger, storageConf)
	store, err := storeManager.Create("rapid-writes")
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, store.Close())
	})

	// 1. Start reader (following)
	rdr, err := store.Reader(ctx, 0, true)
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, rdr.Close())
	})

	count := 5000

	var eg errgroup.Group

	eg.Go(func() error {
		for i := range count {
			if writeErr := store.WriteLine(ctx, fmt.Appendf(nil, "msg %d", i)); writeErr != nil {
				return writeErr
			}
		}

		return nil
	})

	// 3. Read and verify no gaps
	for i := range count {
		line, readErr := rdr.ReadLine(ctx)
		require.NoError(t, readErr)

		require.Equal(t, fmt.Sprintf("msg %d", i), string(line))
	}

	require.NoError(t, eg.Wait())
}

func TestCleanupProbabilities(t *testing.T) {
	t.Parallel()

	// Helper to create a store with specific config
	setupStore := func(t *testing.T, prob float64, limit int) (*sqlitelog.StoreManager, string) {
		ctx := t.Context()
		path := filepath.Join(t.TempDir(), fmt.Sprintf("prob-test-%f.db", prob))

		dbConf := config.Default().Storage.Sqlite
		dbConf.SetPath(path)

		db, err := sqlite.OpenDB(dbConf)
		require.NoError(t, err)

		t.Cleanup(func() { require.NoError(t, db.Close()) })

		logConf := config.Default().Logs.Machine.Storage
		logConf.SetMaxLinesPerMachine(limit)
		logConf.SetCleanupProbability(prob)

		state := state.WrapCore(namespaced.NewState(inmem.Build))

		logger := zaptest.NewLogger(t)
		mgr, err := sqlitelog.NewStoreManager(ctx, db, logConf, state, logger)
		require.NoError(t, err)

		return mgr, fmt.Sprintf("machine-prob-%f", prob)
	}

	tests := []struct {
		name        string
		prob        float64
		limit       int
		writes      int
		expectExact int  // -1 if exact match not expected
		expectLess  bool // if true, expect actual < writes
	}{
		{
			name:        "Probability 0.0 (Never Clean)",
			prob:        0.0,
			limit:       10,
			writes:      50,
			expectExact: 50,
		},
		{
			name:        "Probability 1.0 (Always Clean)",
			prob:        1.0,
			limit:       10,
			writes:      50,
			expectExact: 10,
		},
		{
			name:        "Probability 0.5 (Sometimes Clean)",
			prob:        0.5,
			limit:       10,
			writes:      100,
			expectExact: -1,
			expectLess:  true, // Expect cleanup to have triggered at least once
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			mgr, id := setupStore(t, tt.prob, tt.limit)
			store, err := mgr.Create(id)
			require.NoError(t, err)
			t.Cleanup(func() { require.NoError(t, store.Close()) })

			ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
			t.Cleanup(cancel)

			// Perform Writes
			for i := range tt.writes {
				require.NoError(t, store.WriteLine(ctx, fmt.Appendf(nil, "msg-%d", i)))
			}

			// Check Result
			rdr, err := store.Reader(ctx, -1, false)
			require.NoError(t, err)
			lines := readAllLines(ctx, t, rdr)
			count := len(lines)

			if tt.expectExact != -1 {
				assert.Equal(t, tt.expectExact, count, "Exact count mismatch for prob %f", tt.prob)
			}

			if tt.expectLess {
				assert.Less(t, count, tt.writes, "Cleanup should have triggered at least once")
				assert.GreaterOrEqual(t, count, tt.limit, "Count cannot be less than limit")
			}
		})
	}
}

func TestGlobalSizeCleanup(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)

	// Configure with maxSize=1 (bytes) to force size-based cleanup on every DoCleanup call.
	// cleanupProbability=0 disables per-write cleanup so we isolate the size-based path.
	storageConf := config.Default().Logs.Machine.Storage
	storageConf.SetCleanupProbability(0)
	storageConf.SetMaxSize(1)

	storeManager, omniState := setupDBWithStorageConfig(ctx, t, logger, storageConf)

	machineIDs := []string{"size-m1", "size-m2", "size-m3"}
	linesPerMachine := 1000
	totalRows := len(machineIDs) * linesPerMachine
	msg := []byte("this is a log line used to fill up the database for size-based cleanup testing")

	// Register machines in state so they are not deleted as orphans during DoCleanup.
	for _, id := range machineIDs {
		require.NoError(t, omniState.Create(ctx, omni.NewMachine(id)))
	}

	for _, id := range machineIDs {
		store, err := storeManager.Create(id)
		require.NoError(t, err)

		for range linesPerMachine {
			require.NoError(t, store.WriteLine(ctx, msg))
		}

		require.NoError(t, store.Close())
	}

	// Verify all rows exist before cleanup.
	totalBefore := countTotalRows(ctx, t, storeManager, machineIDs)
	require.Equal(t, totalRows, totalBefore)

	// DoCleanup triggers size-based cleanup since the table is well above maxSize=1.
	// With looping, a single DoCleanup call should delete all rows (3000 total across 3 batches).
	require.NoError(t, storeManager.DoCleanup(ctx))

	totalAfter := countTotalRows(ctx, t, storeManager, machineIDs)

	// The loop should have converged: all 3000 rows deleted, not just a single batch of 1000.
	assert.Equal(t, 0, totalAfter, "size-based cleanup should delete all rows when maxSize=1")
}

// TestGlobalSizeCleanupWithGaps verifies that size-based cleanup converges even when
// the ID space has gaps from per-machine removals.
func TestGlobalSizeCleanupWithGaps(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)

	storageConf := config.Default().Logs.Machine.Storage
	storageConf.SetCleanupProbability(0)
	storageConf.SetMaxSize(1)

	storeManager, omniState := setupDBWithStorageConfig(ctx, t, logger, storageConf)

	msg := []byte("log line for gap test")

	// Write interleaved logs for 3 machines: IDs will be 1,2,3,4,5,6,...
	// Machine assignment: m1 gets IDs 1,4,7,...  m2 gets 2,5,8,...  m3 gets 3,6,9,...
	machineIDs := []string{"gap-m1", "gap-m2", "gap-m3"}
	linesPerMachine := 500

	for _, id := range machineIDs {
		require.NoError(t, omniState.Create(ctx, omni.NewMachine(id)))
	}

	stores := make(map[string]logstore.LogStore, len(machineIDs))
	for _, id := range machineIDs {
		store, err := storeManager.Create(id)
		require.NoError(t, err)

		stores[id] = store
	}

	for range linesPerMachine {
		for _, id := range machineIDs {
			require.NoError(t, stores[id].WriteLine(ctx, msg))
		}
	}

	for _, store := range stores {
		require.NoError(t, store.Close())
	}

	// Remove m2's logs — creates gaps scattered throughout the ID range.
	require.NoError(t, storeManager.Remove(ctx, "gap-m2"))

	// Verify: m1 and m3 each have linesPerMachine, m2 has 0.
	remainingIDs := []string{"gap-m1", "gap-m3"}
	totalBefore := countTotalRows(ctx, t, storeManager, remainingIDs)
	require.Equal(t, 2*linesPerMachine, totalBefore)

	// DoCleanup should delete all remaining rows despite the gaps.
	require.NoError(t, storeManager.DoCleanup(ctx))

	totalAfter := countTotalRows(ctx, t, storeManager, remainingIDs)
	assert.Equal(t, 0, totalAfter, "size-based cleanup should converge even with ID gaps")
}

// countTotalRows returns the total row count across all given machine IDs.
func countTotalRows(ctx context.Context, t *testing.T, mgr *sqlitelog.StoreManager, ids []string) int {
	t.Helper()

	total := 0

	for _, id := range ids {
		store, err := mgr.Create(id)
		require.NoError(t, err)

		rdr, err := store.Reader(ctx, -1, false)
		require.NoError(t, err)

		lines := readAllLines(ctx, t, rdr)
		total += len(lines)

		require.NoError(t, store.Close())
	}

	return total
}

func TestTimeBasedCleanupBatching(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)

	storageConf := config.Default().Logs.Machine.Storage
	storageConf.SetCleanupProbability(0)
	storageConf.SetCleanupOlderThan(24 * time.Hour)

	path := filepath.Join(t.TempDir(), "time-batch.db")
	dbConf := config.Default().Storage.Sqlite
	dbConf.SetPath(path)

	db, err := sqlite.OpenDB(dbConf)
	require.NoError(t, err)
	t.Cleanup(func() { require.NoError(t, db.Close()) })

	omniState := state.WrapCore(namespaced.NewState(inmem.Build))

	var totalCleaned int

	storeManager, err := sqlitelog.NewStoreManager(ctx, db, storageConf, omniState, logger,
		sqlitelog.WithCleanupCallback(func(n int) { totalCleaned += n }),
	)
	require.NoError(t, err)

	machineID := "time-batch-m1"
	require.NoError(t, omniState.Create(ctx, omni.NewMachine(machineID)))

	// Insert rows directly to control timestamps.
	// 1500 old rows (2 days ago) + 200 recent rows (1 hour ago).
	conn, err := db.Take(ctx)
	require.NoError(t, err)

	oldTimestamp := time.Now().Add(-48 * time.Hour).Unix()
	recentTimestamp := time.Now().Add(-1 * time.Hour).Unix()

	const (
		oldRows    = 1500
		recentRows = 200
	)

	for range oldRows {
		q, qErr := sqlitexx.NewQuery(conn, fmt.Sprintf(
			"INSERT INTO %s (%s, %s, %s) VALUES ($mid, $msg, $ts)",
			sqlitelog.TableName, "machine_id", "message", "created_at",
		))
		require.NoError(t, qErr)
		require.NoError(t, q.BindString("$mid", machineID).BindBytes("$msg", []byte("old")).BindInt64("$ts", oldTimestamp).Exec())
	}

	for range recentRows {
		q, qErr := sqlitexx.NewQuery(conn, fmt.Sprintf(
			"INSERT INTO %s (%s, %s, %s) VALUES ($mid, $msg, $ts)",
			sqlitelog.TableName, "machine_id", "message", "created_at",
		))
		require.NoError(t, qErr)
		require.NoError(t, q.BindString("$mid", machineID).BindBytes("$msg", []byte("recent")).BindInt64("$ts", recentTimestamp).Exec())
	}

	db.Put(conn)

	// Verify all rows exist before cleanup.
	totalBefore := countTotalRows(ctx, t, storeManager, []string{machineID})
	require.Equal(t, oldRows+recentRows, totalBefore)

	// DoCleanup should delete all 1500 old rows across multiple batches (1000 + 500).
	require.NoError(t, storeManager.DoCleanup(ctx))

	totalAfter := countTotalRows(ctx, t, storeManager, []string{machineID})
	assert.Equal(t, recentRows, totalAfter, "only recent rows should remain after time-based cleanup")
	assert.Equal(t, oldRows, totalCleaned, "cleanup callback should report all old rows deleted")
}

func TestOrphanLogsCleanup(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second)
	t.Cleanup(cancel)

	logger := zaptest.NewLogger(t)
	storeManager, state := setupDB(ctx, t, logger)

	const (
		numLive        = 1000
		numNonExistent = 1500
	)

	seedLogs := func(prefix string, count int) {
		for i := range count {
			id := fmt.Sprintf("%s-%d", prefix, i)
			store, storeErr := storeManager.Create(id)
			require.NoError(t, storeErr)

			require.NoError(t, store.WriteLine(ctx, fmt.Appendf(nil, "log line A for %s", id)))
			require.NoError(t, store.WriteLine(ctx, fmt.Appendf(nil, "log line B for %s", id)))
			require.NoError(t, store.Close())
		}
	}

	logger.Info("Seed logs...", zap.Int("live", numLive), zap.Int("nonExistent", numNonExistent))

	seedLogs("live", numLive)
	seedLogs("gone", numNonExistent)

	logger.Info("Populate state with machines...")

	for i := range numLive {
		m := omni.NewMachine(fmt.Sprintf("live-%d", i))
		require.NoError(t, state.Create(ctx, m))

		// mark some of the machines as tearing down - their logs should still be preserved
		if i%10 == 0 {
			_, err := state.Teardown(ctx, m.Metadata())
			require.NoError(t, err)
		}
	}

	logger.Info("Trigger cleanup...")

	require.NoError(t, storeManager.DoCleanup(ctx))

	logger.Info("Verify results...")

	checkExistence := func(prefix string, count int, shouldExist bool) {
		for i := range count {
			id := fmt.Sprintf("%s-%d", prefix, i)
			exists, existsErr := storeManager.Exists(ctx, id)
			require.NoError(t, existsErr)

			if shouldExist {
				assert.Truef(t, exists, "Machine %q should have logs", id)
			} else {
				assert.Falsef(t, exists, "Machine %q logs should be deleted", id)
			}
		}
	}

	checkExistence("live", numLive, true)
	checkExistence("gone", numNonExistent, false)
}

func testRead(ctx context.Context, t *testing.T, sqliteStore logstore.LogStore, expectedLines, tailLines int) {
	t.Helper()

	sqliteReader, err := sqliteStore.Reader(ctx, tailLines, false)
	require.NoError(t, err)

	sqliteLines := readAllLines(ctx, t, sqliteReader)

	assert.Len(t, sqliteLines, expectedLines)
}

func readAllLines(ctx context.Context, t *testing.T, rdr logstore.LineReader) []string {
	t.Helper()

	t.Cleanup(func() {
		require.NoError(t, rdr.Close())
	})

	var lines []string

	for {
		line, err := rdr.ReadLine(ctx)
		if err != nil {
			if errors.Is(err, io.EOF) {
				return lines
			}

			require.NoError(t, err)
		}

		lines = append(lines, string(line))
	}
}

func readLines(ctx context.Context, t *testing.T, rdr logstore.LineReader, ch chan<- string) {
	t.Helper()

	t.Cleanup(func() {
		require.NoError(t, rdr.Close())
	})

	for {
		line, err := rdr.ReadLine(ctx)
		if err != nil {
			if errors.Is(err, io.EOF) {
				return
			}

			require.NoError(t, err)
		}

		select {
		case ch <- string(line):
		case <-ctx.Done():
			return
		}
	}
}

func assertLine(ctx context.Context, t *testing.T, lineCh <-chan string, expected string) {
	t.Helper()

	select {
	case line := <-lineCh:
		assert.Equal(t, expected, line)
	case <-ctx.Done():
		require.NoError(t, ctx.Err())
	}
}

// setupDB helper handles the standard SQLite test setup.
func setupDB(ctx context.Context, t *testing.T, logger *zap.Logger) (*sqlitelog.StoreManager, state.State) {
	t.Helper()

	return setupDBWithStorageConfig(ctx, t, logger, config.Default().Logs.Machine.Storage)
}

func setupDBWithStorageConfig(ctx context.Context, t *testing.T, logger *zap.Logger, storageConf config.LogsMachineStorage) (*sqlitelog.StoreManager, state.State) {
	t.Helper()

	path := filepath.Join(t.TempDir(), "test.db")
	conf := config.Default().Storage.Sqlite
	conf.SetPath(path)

	db, err := sqlite.OpenDB(conf)
	require.NoError(t, err)

	t.Cleanup(func() {
		require.NoError(t, db.Close())
	})

	state := state.WrapCore(namespaced.NewState(inmem.Build))

	storeManager, err := sqlitelog.NewStoreManager(ctx, db, storageConf, state, logger)
	require.NoError(t, err)

	return storeManager, state
}
