package memory

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/sipeed/picoclaw/pkg/providers"
)

func newTestStore(t *testing.T) *JSONLStore {
	t.Helper()
	store, err := NewJSONLStore(t.TempDir())
	if err != nil {
		t.Fatalf("NewJSONLStore: %v", err)
	}
	return store
}

func TestNewJSONLStore_CreatesDirectory(t *testing.T) {
	dir := filepath.Join(t.TempDir(), "nested", "sessions")
	store, err := NewJSONLStore(dir)
	if err != nil {
		t.Fatalf("NewJSONLStore: %v", err)
	}
	defer store.Close()

	info, err := os.Stat(dir)
	if err != nil {
		t.Fatalf("Stat: %v", err)
	}
	if !info.IsDir() {
		t.Errorf("expected directory, got file")
	}
}

func TestAddMessage_BasicRoundtrip(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	err := store.AddMessage(ctx, "s1", "user", "hello")
	if err != nil {
		t.Fatalf("AddMessage: %v", err)
	}
	err = store.AddMessage(ctx, "s1", "assistant", "hi there")
	if err != nil {
		t.Fatalf("AddMessage: %v", err)
	}

	history, err := store.GetHistory(ctx, "s1")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 2 {
		t.Fatalf("expected 2 messages, got %d", len(history))
	}
	if history[0].Role != "user" || history[0].Content != "hello" {
		t.Errorf("msg[0] = %+v", history[0])
	}
	if history[1].Role != "assistant" || history[1].Content != "hi there" {
		t.Errorf("msg[1] = %+v", history[1])
	}
}

func TestAddMessage_AutoCreatesSession(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	// Adding a message to a non-existent session should work.
	err := store.AddMessage(ctx, "new-session", "user", "first message")
	if err != nil {
		t.Fatalf("AddMessage: %v", err)
	}

	history, err := store.GetHistory(ctx, "new-session")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 1 {
		t.Fatalf("expected 1 message, got %d", len(history))
	}
}

func TestAddFullMessage_WithToolCalls(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	msg := providers.Message{
		Role:    "assistant",
		Content: "Let me search that.",
		ToolCalls: []providers.ToolCall{
			{
				ID:   "call_abc",
				Type: "function",
				Function: &providers.FunctionCall{
					Name:      "web_search",
					Arguments: `{"q":"golang jsonl"}`,
				},
			},
		},
	}

	err := store.AddFullMessage(ctx, "tc", msg)
	if err != nil {
		t.Fatalf("AddFullMessage: %v", err)
	}

	history, err := store.GetHistory(ctx, "tc")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 1 {
		t.Fatalf("expected 1, got %d", len(history))
	}
	if len(history[0].ToolCalls) != 1 {
		t.Fatalf("expected 1 tool call, got %d", len(history[0].ToolCalls))
	}
	tc := history[0].ToolCalls[0]
	if tc.ID != "call_abc" {
		t.Errorf("tool call ID = %q", tc.ID)
	}
	if tc.Function == nil || tc.Function.Name != "web_search" {
		t.Errorf("tool call function = %+v", tc.Function)
	}
}

func TestAddFullMessage_PreservesModelName(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	msg := providers.Message{
		Role:      "assistant",
		Content:   "done",
		ModelName: "gpt-5.4-mini",
	}

	if err := store.AddFullMessage(ctx, "model-name", msg); err != nil {
		t.Fatalf("AddFullMessage: %v", err)
	}

	history, err := store.GetHistory(ctx, "model-name")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 1 {
		t.Fatalf("expected 1, got %d", len(history))
	}
	if history[0].ModelName != "gpt-5.4-mini" {
		t.Fatalf("ModelName = %q, want %q", history[0].ModelName, "gpt-5.4-mini")
	}
}

func TestAddFullMessage_ToolCallID(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	msg := providers.Message{
		Role:       "tool",
		Content:    "search results here",
		ToolCallID: "call_abc",
	}

	err := store.AddFullMessage(ctx, "tr", msg)
	if err != nil {
		t.Fatalf("AddFullMessage: %v", err)
	}

	history, err := store.GetHistory(ctx, "tr")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 1 {
		t.Fatalf("expected 1, got %d", len(history))
	}
	if history[0].ToolCallID != "call_abc" {
		t.Errorf("ToolCallID = %q", history[0].ToolCallID)
	}
}

func TestAddFullMessage_DropsTransientAssistantThought(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	err := store.AddFullMessage(ctx, "transient-thought", providers.Message{
		Role:             "assistant",
		ReasoningContent: "internal chain of thought",
	})
	if err != nil {
		t.Fatalf("AddFullMessage: %v", err)
	}

	history, err := store.GetHistory(ctx, "transient-thought")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 0 {
		t.Fatalf("expected transient thought to be discarded, got %d messages", len(history))
	}
}

func TestGetHistory_EmptySession(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	history, err := store.GetHistory(ctx, "nonexistent")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if history == nil {
		t.Fatal("expected non-nil empty slice")
	}
	if len(history) != 0 {
		t.Errorf("expected 0 messages, got %d", len(history))
	}
}

func TestGetHistory_Ordering(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	for i := 0; i < 5; i++ {
		err := store.AddMessage(
			ctx, "order",
			"user",
			string(rune('a'+i)),
		)
		if err != nil {
			t.Fatalf("AddMessage(%d): %v", i, err)
		}
	}

	history, err := store.GetHistory(ctx, "order")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 5 {
		t.Fatalf("expected 5, got %d", len(history))
	}
	for i := 0; i < 5; i++ {
		expected := string(rune('a' + i))
		if history[i].Content != expected {
			t.Errorf("msg[%d].Content = %q, want %q", i, history[i].Content, expected)
		}
	}
}

func TestSetSummary_GetSummary(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	// No summary yet.
	summary, err := store.GetSummary(ctx, "s1")
	if err != nil {
		t.Fatalf("GetSummary: %v", err)
	}
	if summary != "" {
		t.Errorf("expected empty, got %q", summary)
	}

	// Set a summary.
	err = store.SetSummary(ctx, "s1", "talked about Go")
	if err != nil {
		t.Fatalf("SetSummary: %v", err)
	}

	summary, err = store.GetSummary(ctx, "s1")
	if err != nil {
		t.Fatalf("GetSummary: %v", err)
	}
	if summary != "talked about Go" {
		t.Errorf("summary = %q", summary)
	}

	// Update summary.
	err = store.SetSummary(ctx, "s1", "updated summary")
	if err != nil {
		t.Fatalf("SetSummary: %v", err)
	}

	summary, err = store.GetSummary(ctx, "s1")
	if err != nil {
		t.Fatalf("GetSummary: %v", err)
	}
	if summary != "updated summary" {
		t.Errorf("summary = %q", summary)
	}
}

func TestSetHistory_DropsTransientAssistantThought(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	newHistory := []providers.Message{
		{Role: "user", Content: "hello"},
		{Role: "assistant", ReasoningContent: "internal chain of thought"},
		{Role: "assistant", Content: "visible answer", ReasoningContent: "visible thought"},
	}

	err := store.SetHistory(ctx, "replace", newHistory)
	if err != nil {
		t.Fatalf("SetHistory: %v", err)
	}

	history, err := store.GetHistory(ctx, "replace")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 2 {
		t.Fatalf("expected transient thought to be removed, got %d messages", len(history))
	}
	if history[0].Role != "user" || history[0].Content != "hello" {
		t.Fatalf("history[0] = %+v, want user/hello", history[0])
	}
	if history[1].Role != "assistant" || history[1].Content != "visible answer" ||
		history[1].ReasoningContent != "visible thought" {
		t.Fatalf("history[1] = %+v, want assistant visible answer with reasoning", history[1])
	}

	data, err := os.ReadFile(store.jsonlPath("replace"))
	if err != nil {
		t.Fatalf("ReadFile(jsonl): %v", err)
	}
	lines := strings.Split(strings.TrimSpace(string(data)), "\n")
	if len(lines) != 2 {
		t.Fatalf("jsonl line count = %d, want 2", len(lines))
	}
}

func TestSessionMetaScopeAndAliasesPersist(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	scope := json.RawMessage(`{"version":1,"channel":"telegram","values":{"chat":"group:c1"}}`)
	aliases := []string{"legacy:one", "legacy:one", "canonical"}
	if err := store.UpsertSessionMeta(ctx, "canonical", scope, aliases); err != nil {
		t.Fatalf("UpsertSessionMeta() error = %v", err)
	}

	meta, err := store.GetSessionMeta(ctx, "canonical")
	if err != nil {
		t.Fatalf("GetSessionMeta() error = %v", err)
	}
	var gotScope map[string]any
	if err := json.Unmarshal(meta.Scope, &gotScope); err != nil {
		t.Fatalf("Unmarshal(meta.Scope) error = %v", err)
	}
	var wantScope map[string]any
	if err := json.Unmarshal(scope, &wantScope); err != nil {
		t.Fatalf("Unmarshal(scope) error = %v", err)
	}
	if !reflect.DeepEqual(gotScope, wantScope) {
		t.Fatalf("meta.Scope = %#v, want %#v", gotScope, wantScope)
	}
	if len(meta.Aliases) != 1 || meta.Aliases[0] != "legacy:one" {
		t.Fatalf("meta.Aliases = %#v, want [legacy:one]", meta.Aliases)
	}
}

func TestResolveSessionKeyByAlias(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	if err := store.AddMessage(ctx, "canonical", "user", "hello"); err != nil {
		t.Fatalf("AddMessage() error = %v", err)
	}
	if err := store.UpsertSessionMeta(ctx, "canonical", nil, []string{"legacy:key"}); err != nil {
		t.Fatalf("UpsertSessionMeta() error = %v", err)
	}

	resolved, found, err := store.ResolveSessionKey(ctx, "legacy:key")
	if err != nil {
		t.Fatalf("ResolveSessionKey() error = %v", err)
	}
	if !found {
		t.Fatal("ResolveSessionKey() did not find alias")
	}
	if resolved != "canonical" {
		t.Fatalf("resolved = %q, want %q", resolved, "canonical")
	}
}

func TestResolveSessionKeyByAlias_PrefersMetadataOverLegacyFile(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	if err := store.AddMessage(ctx, "legacy:key", "user", "legacy"); err != nil {
		t.Fatalf("AddMessage(legacy) error = %v", err)
	}
	if err := store.AddMessage(ctx, "canonical", "user", "canonical"); err != nil {
		t.Fatalf("AddMessage(canonical) error = %v", err)
	}
	if err := store.UpsertSessionMeta(ctx, "canonical", nil, []string{"legacy:key"}); err != nil {
		t.Fatalf("UpsertSessionMeta() error = %v", err)
	}

	resolved, found, err := store.ResolveSessionKey(ctx, "legacy:key")
	if err != nil {
		t.Fatalf("ResolveSessionKey() error = %v", err)
	}
	if !found {
		t.Fatal("ResolveSessionKey() did not find alias")
	}
	if resolved != "canonical" {
		t.Fatalf("resolved = %q, want %q", resolved, "canonical")
	}
}

func TestResolveSessionKey_DirectHitSkipsCorruptMetadata(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	if err := store.AddMessage(ctx, "canonical", "user", "hello"); err != nil {
		t.Fatalf("AddMessage() error = %v", err)
	}
	if err := os.WriteFile(
		filepath.Join(store.dir, "broken.meta.json"),
		[]byte("{not-json"),
		0o644,
	); err != nil {
		t.Fatalf("WriteFile(broken.meta.json) error = %v", err)
	}

	resolved, found, err := store.ResolveSessionKey(ctx, "canonical")
	if err != nil {
		t.Fatalf("ResolveSessionKey() error = %v", err)
	}
	if !found {
		t.Fatal("ResolveSessionKey() did not find direct session")
	}
	if resolved != "canonical" {
		t.Fatalf("resolved = %q, want %q", resolved, "canonical")
	}
}

func TestResolveSessionKey_SkipsCorruptMetadataDuringAliasScan(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	if err := store.AddMessage(ctx, "canonical", "user", "hello"); err != nil {
		t.Fatalf("AddMessage() error = %v", err)
	}
	if err := store.UpsertSessionMeta(ctx, "canonical", nil, []string{"legacy:key"}); err != nil {
		t.Fatalf("UpsertSessionMeta() error = %v", err)
	}
	if err := os.WriteFile(
		filepath.Join(store.dir, "broken.meta.json"),
		[]byte("{not-json"),
		0o644,
	); err != nil {
		t.Fatalf("WriteFile(broken.meta.json) error = %v", err)
	}

	resolved, found, err := store.ResolveSessionKey(ctx, "legacy:key")
	if err != nil {
		t.Fatalf("ResolveSessionKey() error = %v", err)
	}
	if !found {
		t.Fatal("ResolveSessionKey() did not find alias")
	}
	if resolved != "canonical" {
		t.Fatalf("resolved = %q, want %q", resolved, "canonical")
	}
}

func TestTruncateHistory_KeepLast(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	for i := 0; i < 10; i++ {
		err := store.AddMessage(
			ctx, "trunc",
			"user",
			string(rune('a'+i)),
		)
		if err != nil {
			t.Fatalf("AddMessage: %v", err)
		}
	}

	err := store.TruncateHistory(ctx, "trunc", 4)
	if err != nil {
		t.Fatalf("TruncateHistory: %v", err)
	}

	history, err := store.GetHistory(ctx, "trunc")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 4 {
		t.Fatalf("expected 4, got %d", len(history))
	}
	// Should be the last 4: g, h, i, j
	if history[0].Content != "g" {
		t.Errorf("first kept = %q, want 'g'", history[0].Content)
	}
	if history[3].Content != "j" {
		t.Errorf("last kept = %q, want 'j'", history[3].Content)
	}
}

func TestTruncateHistory_KeepZero(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	for i := 0; i < 5; i++ {
		err := store.AddMessage(ctx, "empty", "user", "msg")
		if err != nil {
			t.Fatalf("AddMessage: %v", err)
		}
	}

	err := store.TruncateHistory(ctx, "empty", 0)
	if err != nil {
		t.Fatalf("TruncateHistory: %v", err)
	}

	history, err := store.GetHistory(ctx, "empty")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 0 {
		t.Errorf("expected 0, got %d", len(history))
	}
}

func TestTruncateHistory_KeepMoreThanExists(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	for i := 0; i < 3; i++ {
		err := store.AddMessage(ctx, "few", "user", "msg")
		if err != nil {
			t.Fatalf("AddMessage: %v", err)
		}
	}

	// Keep 100, but only 3 exist — should keep all.
	err := store.TruncateHistory(ctx, "few", 100)
	if err != nil {
		t.Fatalf("TruncateHistory: %v", err)
	}

	history, err := store.GetHistory(ctx, "few")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 3 {
		t.Errorf("expected 3, got %d", len(history))
	}
}

func TestSetHistory_ReplacesAll(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	// Add some initial messages.
	for i := 0; i < 5; i++ {
		err := store.AddMessage(ctx, "replace", "user", "old")
		if err != nil {
			t.Fatalf("AddMessage: %v", err)
		}
	}

	// Replace with new history.
	newHistory := []providers.Message{
		{Role: "user", Content: "new1"},
		{Role: "assistant", Content: "new2"},
	}
	err := store.SetHistory(ctx, "replace", newHistory)
	if err != nil {
		t.Fatalf("SetHistory: %v", err)
	}

	history, err := store.GetHistory(ctx, "replace")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 2 {
		t.Fatalf("expected 2, got %d", len(history))
	}
	if history[0].Content != "new1" || history[1].Content != "new2" {
		t.Errorf("history = %+v", history)
	}
}

func TestSetHistory_ResetsSkip(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	// Add messages and truncate.
	for i := 0; i < 10; i++ {
		err := store.AddMessage(ctx, "skip-reset", "user", "old")
		if err != nil {
			t.Fatalf("AddMessage: %v", err)
		}
	}
	err := store.TruncateHistory(ctx, "skip-reset", 3)
	if err != nil {
		t.Fatalf("TruncateHistory: %v", err)
	}

	// SetHistory should reset skip to 0.
	newHistory := []providers.Message{
		{Role: "user", Content: "fresh"},
	}
	err = store.SetHistory(ctx, "skip-reset", newHistory)
	if err != nil {
		t.Fatalf("SetHistory: %v", err)
	}

	history, err := store.GetHistory(ctx, "skip-reset")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 1 {
		t.Fatalf("expected 1, got %d", len(history))
	}
	if history[0].Content != "fresh" {
		t.Errorf("content = %q", history[0].Content)
	}
}

func TestColonInKey(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	err := store.AddMessage(ctx, "telegram:123", "user", "hi")
	if err != nil {
		t.Fatalf("AddMessage: %v", err)
	}

	history, err := store.GetHistory(ctx, "telegram:123")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 1 {
		t.Fatalf("expected 1, got %d", len(history))
	}

	// Verify the file is named with underscore.
	jsonlFile := filepath.Join(store.dir, "telegram_123.jsonl")
	if _, statErr := os.Stat(jsonlFile); statErr != nil {
		t.Errorf("expected file %s to exist: %v", jsonlFile, statErr)
	}
}

func TestCompact_RemovesSkippedMessages(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	// Write 10 messages, then truncate to keep last 3.
	for i := 0; i < 10; i++ {
		err := store.AddMessage(ctx, "compact", "user", string(rune('a'+i)))
		if err != nil {
			t.Fatalf("AddMessage: %v", err)
		}
	}
	err := store.TruncateHistory(ctx, "compact", 3)
	if err != nil {
		t.Fatalf("TruncateHistory: %v", err)
	}

	// Before compact: file still has 10 lines.
	allOnDisk, err := readMessages(store.jsonlPath("compact"), 0)
	if err != nil {
		t.Fatalf("readMessages: %v", err)
	}
	if len(allOnDisk) != 10 {
		t.Fatalf("before compact: expected 10 on disk, got %d", len(allOnDisk))
	}

	// Compact.
	err = store.Compact(ctx, "compact")
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}

	// After compact: file should have only 3 lines.
	allOnDisk, err = readMessages(store.jsonlPath("compact"), 0)
	if err != nil {
		t.Fatalf("readMessages: %v", err)
	}
	if len(allOnDisk) != 3 {
		t.Fatalf("after compact: expected 3 on disk, got %d", len(allOnDisk))
	}

	// GetHistory should still return the same 3 messages.
	history, err := store.GetHistory(ctx, "compact")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 3 {
		t.Fatalf("expected 3, got %d", len(history))
	}
	if history[0].Content != "h" || history[2].Content != "j" {
		t.Errorf("wrong content: %+v", history)
	}
}

func TestCompact_NoOpWhenNoSkip(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	for i := 0; i < 5; i++ {
		err := store.AddMessage(ctx, "noop", "user", "msg")
		if err != nil {
			t.Fatalf("AddMessage: %v", err)
		}
	}

	// Compact without prior truncation — should be a no-op.
	err := store.Compact(ctx, "noop")
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}

	history, err := store.GetHistory(ctx, "noop")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 5 {
		t.Errorf("expected 5, got %d", len(history))
	}
}

func TestCompact_ThenAppend(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	for i := 0; i < 8; i++ {
		err := store.AddMessage(ctx, "cap", "user", string(rune('a'+i)))
		if err != nil {
			t.Fatalf("AddMessage: %v", err)
		}
	}

	err := store.TruncateHistory(ctx, "cap", 2)
	if err != nil {
		t.Fatalf("TruncateHistory: %v", err)
	}
	err = store.Compact(ctx, "cap")
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}

	// Append after compaction should work correctly.
	err = store.AddMessage(ctx, "cap", "user", "new")
	if err != nil {
		t.Fatalf("AddMessage after compact: %v", err)
	}

	history, err := store.GetHistory(ctx, "cap")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 3 {
		t.Fatalf("expected 3, got %d", len(history))
	}
	// g, h (kept from truncation), new (appended after compaction).
	if history[0].Content != "g" {
		t.Errorf("first = %q, want 'g'", history[0].Content)
	}
	if history[2].Content != "new" {
		t.Errorf("last = %q, want 'new'", history[2].Content)
	}
}

func TestTruncateHistory_StaleMetaCount(t *testing.T) {
	// Simulates a crash between JSONL append and meta update in addMsg:
	// file has N+1 lines but meta.Count is still N. TruncateHistory must
	// reconcile with the real line count so that keepLast is accurate.
	store := newTestStore(t)
	ctx := context.Background()

	// Write 10 messages normally (meta.Count = 10).
	for i := 0; i < 10; i++ {
		err := store.AddMessage(ctx, "stale", "user", string(rune('a'+i)))
		if err != nil {
			t.Fatalf("AddMessage: %v", err)
		}
	}

	// Simulate crash: append a line to JSONL but do NOT update meta.
	// This leaves meta.Count = 10 while the file has 11 lines.
	jsonlPath := store.jsonlPath("stale")
	f, err := os.OpenFile(jsonlPath, os.O_WRONLY|os.O_APPEND, 0o644)
	if err != nil {
		t.Fatalf("open for append: %v", err)
	}
	_, err = f.WriteString(`{"role":"user","content":"orphan"}` + "\n")
	if err != nil {
		t.Fatalf("write orphan: %v", err)
	}
	f.Close()

	// TruncateHistory(keepLast=4) should keep the last 4 of 11 lines,
	// not the last 4 of 10.
	err = store.TruncateHistory(ctx, "stale", 4)
	if err != nil {
		t.Fatalf("TruncateHistory: %v", err)
	}

	history, err := store.GetHistory(ctx, "stale")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 4 {
		t.Fatalf("expected 4, got %d", len(history))
	}
	// Last 4 of [a,b,c,d,e,f,g,h,i,j,orphan] = [h,i,j,orphan]
	if history[0].Content != "h" {
		t.Errorf("first kept = %q, want 'h'", history[0].Content)
	}
	if history[3].Content != "orphan" {
		t.Errorf("last kept = %q, want 'orphan'", history[3].Content)
	}
}

func TestTruncateHistory_IgnoresTransientThoughtForKeepLast(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()
	sessionKey := "transient-keep-last"
	now := time.Now()

	rawJSONL := strings.Join([]string{
		`{"role":"user","content":"a"}`,
		`{"role":"assistant","content":"b"}`,
		`{"role":"assistant","content":"","reasoning_content":"dangling thought"}`,
		`{"role":"user","content":"c"}`,
		`{"role":"assistant","content":"d"}`,
	}, "\n") + "\n"
	if err := os.WriteFile(store.jsonlPath(sessionKey), []byte(rawJSONL), 0o644); err != nil {
		t.Fatalf("WriteFile(jsonl): %v", err)
	}
	if err := store.writeMeta(sessionKey, SessionMeta{
		Key:       sessionKey,
		Count:     5,
		Skip:      0,
		CreatedAt: now,
		UpdatedAt: now,
	}); err != nil {
		t.Fatalf("writeMeta: %v", err)
	}

	if err := store.TruncateHistory(ctx, sessionKey, 2); err != nil {
		t.Fatalf("TruncateHistory: %v", err)
	}

	history, err := store.GetHistory(ctx, sessionKey)
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 2 {
		t.Fatalf("expected 2 retained messages, got %d", len(history))
	}
	if history[0].Content != "c" || history[1].Content != "d" {
		t.Fatalf("kept history = %+v, want c,d", history)
	}

	meta, err := store.readMeta(sessionKey)
	if err != nil {
		t.Fatalf("readMeta: %v", err)
	}
	if meta.Skip != 2 {
		t.Fatalf("meta.Skip = %d, want 2 raw lines skipped", meta.Skip)
	}
}

func TestCrashRecovery_PartialLine(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	// Write a valid message first.
	err := store.AddMessage(ctx, "crash", "user", "valid")
	if err != nil {
		t.Fatalf("AddMessage: %v", err)
	}

	// Simulate a crash by appending a partial JSON line directly.
	jsonlPath := store.jsonlPath("crash")
	f, err := os.OpenFile(jsonlPath, os.O_WRONLY|os.O_APPEND, 0o644)
	if err != nil {
		t.Fatalf("open for append: %v", err)
	}
	_, err = f.WriteString(`{"role":"user","content":"incomple`)
	if err != nil {
		t.Fatalf("write partial: %v", err)
	}
	f.Close()

	// GetHistory should return only the valid message.
	history, err := store.GetHistory(ctx, "crash")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 1 {
		t.Fatalf("expected 1 valid message, got %d", len(history))
	}
	if history[0].Content != "valid" {
		t.Errorf("content = %q", history[0].Content)
	}
}

func TestPersistence_AcrossInstances(t *testing.T) {
	dir := t.TempDir()
	ctx := context.Background()

	// Write with first instance.
	store1, err := NewJSONLStore(dir)
	if err != nil {
		t.Fatalf("NewJSONLStore: %v", err)
	}
	err = store1.AddMessage(ctx, "persist", "user", "remember me")
	if err != nil {
		t.Fatalf("AddMessage: %v", err)
	}
	err = store1.SetSummary(ctx, "persist", "a test session")
	if err != nil {
		t.Fatalf("SetSummary: %v", err)
	}
	store1.Close()

	// Read with second instance.
	store2, err := NewJSONLStore(dir)
	if err != nil {
		t.Fatalf("NewJSONLStore: %v", err)
	}
	defer store2.Close()

	history, err := store2.GetHistory(ctx, "persist")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	if len(history) != 1 || history[0].Content != "remember me" {
		t.Errorf("history = %+v", history)
	}

	summary, err := store2.GetSummary(ctx, "persist")
	if err != nil {
		t.Fatalf("GetSummary: %v", err)
	}
	if summary != "a test session" {
		t.Errorf("summary = %q", summary)
	}
}

func TestConcurrent_AddAndRead(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	var wg sync.WaitGroup
	const goroutines = 10
	const msgsPerGoroutine = 20

	// Concurrent writes.
	for g := 0; g < goroutines; g++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < msgsPerGoroutine; i++ {
				_ = store.AddMessage(ctx, "concurrent", "user", "msg")
			}
		}()
	}
	wg.Wait()

	history, err := store.GetHistory(ctx, "concurrent")
	if err != nil {
		t.Fatalf("GetHistory: %v", err)
	}
	expected := goroutines * msgsPerGoroutine
	if len(history) != expected {
		t.Errorf("expected %d messages, got %d", expected, len(history))
	}
}

func TestConcurrent_SummarizeRace(t *testing.T) {
	// Simulates the #704 race: one goroutine adds messages while
	// another truncates + sets summary — like summarizeSession().
	store := newTestStore(t)
	ctx := context.Background()

	// Seed with some messages.
	for i := 0; i < 20; i++ {
		err := store.AddMessage(ctx, "race", "user", "seed")
		if err != nil {
			t.Fatalf("AddMessage: %v", err)
		}
	}

	var wg sync.WaitGroup

	// Writer goroutine (main agent loop).
	wg.Add(1)
	go func() {
		defer wg.Done()
		for i := 0; i < 50; i++ {
			_ = store.AddMessage(ctx, "race", "user", "new")
		}
	}()

	// Summarizer goroutine (background task).
	wg.Add(1)
	go func() {
		defer wg.Done()
		for i := 0; i < 10; i++ {
			_ = store.SetSummary(ctx, "race", "summary")
			_ = store.TruncateHistory(ctx, "race", 5)
		}
	}()

	wg.Wait()

	// Verify the store is still in a consistent state.
	_, err := store.GetHistory(ctx, "race")
	if err != nil {
		t.Fatalf("GetHistory after race: %v", err)
	}
	_, err = store.GetSummary(ctx, "race")
	if err != nil {
		t.Fatalf("GetSummary after race: %v", err)
	}
}

func TestMultipleSessions_Isolation(t *testing.T) {
	store := newTestStore(t)
	ctx := context.Background()

	err := store.AddMessage(ctx, "s1", "user", "msg for s1")
	if err != nil {
		t.Fatalf("AddMessage: %v", err)
	}
	err = store.AddMessage(ctx, "s2", "user", "msg for s2")
	if err != nil {
		t.Fatalf("AddMessage: %v", err)
	}

	h1, err := store.GetHistory(ctx, "s1")
	if err != nil {
		t.Fatalf("GetHistory s1: %v", err)
	}
	h2, err := store.GetHistory(ctx, "s2")
	if err != nil {
		t.Fatalf("GetHistory s2: %v", err)
	}

	if len(h1) != 1 || h1[0].Content != "msg for s1" {
		t.Errorf("s1 history = %+v", h1)
	}
	if len(h2) != 1 || h2[0].Content != "msg for s2" {
		t.Errorf("s2 history = %+v", h2)
	}
}

func TestStore_SetsCreatedAtWhenNil(t *testing.T) {
	type writeOp struct {
		name string
		fn   func(store *JSONLStore, key string) (expectedCount int)
	}

	ops := []writeOp{
		{
			name: "AddMessage",
			fn: func(store *JSONLStore, key string) int {
				if err := store.AddMessage(context.Background(), key, "user", "hello"); err != nil {
					t.Fatalf("AddMessage: %v", err)
				}
				return 1
			},
		},
		{
			name: "AddFullMessage",
			fn: func(store *JSONLStore, key string) int {
				if err := store.AddFullMessage(context.Background(), key, providers.Message{
					Role:    "user",
					Content: "hello from full",
				}); err != nil {
					t.Fatalf("AddFullMessage: %v", err)
				}
				return 1
			},
		},
		{
			name: "SetHistory",
			fn: func(store *JSONLStore, key string) int {
				if err := store.SetHistory(context.Background(), key, []providers.Message{
					{Role: "user", Content: "msg1"},
					{Role: "assistant", Content: "msg2"},
				}); err != nil {
					t.Fatalf("SetHistory: %v", err)
				}
				return 2
			},
		},
	}

	for _, op := range ops {
		t.Run(op.name, func(t *testing.T) {
			store := newTestStore(t)
			key := "s1"

			before := time.Now().Add(-time.Second)
			expectedCount := op.fn(store, key)
			after := time.Now().Add(time.Second)

			history, err := store.GetHistory(context.Background(), key)
			if err != nil {
				t.Fatalf("GetHistory: %v", err)
			}
			if len(history) != expectedCount {
				t.Fatalf("expected %d messages, got %d", expectedCount, len(history))
			}
			for i := range history {
				if history[i].CreatedAt == nil || history[i].CreatedAt.IsZero() {
					t.Errorf("message %d CreatedAt is zero — not set by %s", i, op.name)
				}
				if history[i].CreatedAt.Before(before) || history[i].CreatedAt.After(after) {
					t.Errorf(
						"message %d CreatedAt %v outside expected window [%v, %v]",
						i, history[i].CreatedAt, before, after,
					)
				}
			}
		})
	}
}

func TestStore_PreservesExistingCreatedAt(t *testing.T) {
	t1 := time.Date(2026, 1, 1, 10, 0, 0, 0, time.UTC)
	t2 := time.Date(2026, 1, 1, 11, 0, 0, 0, time.UTC)

	type writeOp struct {
		name      string
		fn        func(store *JSONLStore, key string)
		wantTimes []time.Time
	}

	ops := []writeOp{
		{
			name: "AddFullMessage",
			fn: func(store *JSONLStore, key string) {
				if err := store.AddFullMessage(context.Background(), key, providers.Message{
					Role:      "user",
					Content:   "custom time",
					CreatedAt: &t1,
				}); err != nil {
					t.Fatalf("AddFullMessage: %v", err)
				}
			},
			wantTimes: []time.Time{t1},
		},
		{
			name: "SetHistory",
			fn: func(store *JSONLStore, key string) {
				if err := store.SetHistory(context.Background(), key, []providers.Message{
					{Role: "user", Content: "msg1", CreatedAt: &t1},
					{Role: "assistant", Content: "msg2", CreatedAt: &t2},
				}); err != nil {
					t.Fatalf("SetHistory: %v", err)
				}
			},
			wantTimes: []time.Time{t1, t2},
		},
	}

	for _, op := range ops {
		t.Run(op.name, func(t *testing.T) {
			store := newTestStore(t)
			key := "s1"

			op.fn(store, key)

			history, err := store.GetHistory(context.Background(), key)
			if err != nil {
				t.Fatalf("GetHistory: %v", err)
			}
			if len(history) != len(op.wantTimes) {
				t.Fatalf("expected %d messages, got %d", len(op.wantTimes), len(history))
			}
			for i, want := range op.wantTimes {
				if history[i].CreatedAt == nil || !history[i].CreatedAt.Equal(want) {
					t.Errorf(
						"message %d CreatedAt = %v, want %v (should preserve caller-provided time)",
						i, history[i].CreatedAt, want,
					)
				}
			}
		})
	}
}

func BenchmarkAddMessage(b *testing.B) {
	dir := b.TempDir()
	store, err := NewJSONLStore(dir)
	if err != nil {
		b.Fatalf("NewJSONLStore: %v", err)
	}
	defer store.Close()
	ctx := context.Background()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = store.AddMessage(ctx, "bench", "user", "benchmark message content")
	}
}

func BenchmarkGetHistory_100(b *testing.B) {
	dir := b.TempDir()
	store, err := NewJSONLStore(dir)
	if err != nil {
		b.Fatalf("NewJSONLStore: %v", err)
	}
	defer store.Close()
	ctx := context.Background()

	for i := 0; i < 100; i++ {
		_ = store.AddMessage(ctx, "bench", "user", "message content")
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _ = store.GetHistory(ctx, "bench")
	}
}

func BenchmarkGetHistory_1000(b *testing.B) {
	dir := b.TempDir()
	store, err := NewJSONLStore(dir)
	if err != nil {
		b.Fatalf("NewJSONLStore: %v", err)
	}
	defer store.Close()
	ctx := context.Background()

	for i := 0; i < 1000; i++ {
		_ = store.AddMessage(ctx, "bench", "user", "message content")
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _ = store.GetHistory(ctx, "bench")
	}
}