Agent-System/internal/ctxmgr/chat_window_test.go

package ctxmgr_test

import (
	"context"
	"fmt"
	"os"
	"testing"
	"time"

	"chat-agent-system/internal/ctxmgr"

	"github.com/jackc/pgx/v5/pgxpool"
)

func TestDeriveMessageCategory(t *testing.T) {
	tests := []struct {
		name     string
		role     string
		intent   string
		branch   string
		meta     map[string]any
		expected string
	}{
		{
			name:     "milestone metadata",
			role:     "user",
			meta:     map[string]any{"milestone": true},
			expected: ctxmgr.CategoryMilestone,
		},
		{
			name:     "attachments file op",
			role:     "user",
			meta:     map[string]any{"attachments": []any{"file.pdf"}},
			expected: ctxmgr.CategoryFile,
		},
		{
			name:     "file_op export",
			role:     "assistant",
			meta:     map[string]any{"file_op": "export"},
			expected: ctxmgr.CategoryFile,
		},
		{
			name:     "interactive_card review",
			role:     "assistant",
			meta:     map[string]any{"interactive_card": map[string]any{"type": "confirm_reject"}},
			expected: ctxmgr.CategoryReview,
		},
		{
			name:     "task_request task",
			role:     "assistant",
			meta:     map[string]any{"task_request": map[string]any{"taskKey": "intake"}},
			expected: ctxmgr.CategoryTask,
		},
		{
			name:     "DOCGOV intent review",
			role:     "user",
			intent:   "DOCGOV_REVIEW",
			expected: ctxmgr.CategoryReview,
		},
		{
			name:     "TRIGGER_TASK intent task",
			role:     "user",
			intent:   "TRIGGER_TASK",
			expected: ctxmgr.CategoryTask,
		},
		{
			name:     "QUERY_PROGRESS consult",
			role:     "user",
			intent:   "QUERY_PROGRESS",
			expected: ctxmgr.CategoryConsult,
		},
		{
			name:     "CASUAL_CHAT casual",
			role:     "user",
			intent:   "CASUAL_CHAT",
			expected: ctxmgr.CategoryCasual,
		},
		{
			name:     "system role system",
			role:     "system",
			expected: ctxmgr.CategorySystem,
		},
		{
			name:     "fallback uncategorized",
			role:     "user",
			intent:   "",
			expected: ctxmgr.CategoryUncategorized,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			actual := ctxmgr.DeriveMessageCategory(tt.role, tt.intent, tt.branch, tt.meta)
			if actual != tt.expected {
				t.Errorf("DeriveMessageCategory() = %s, want %s", actual, tt.expected)
			}
		})
	}
}

func TestChatWindow_MemoryEquivalenceAndPagination(t *testing.T) {
	mgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.DefaultTimeDecayConfig())
	ctx := context.Background()
	sessID := "sess_test_window_mem"

	// 1. 参数校验
	_, err := mgr.GetChatWindow(ctx, ctxmgr.ChatWindowQuery{SessionID: ""})
	if err == nil {
		t.Errorf("expected error for empty sessionID")
	}

	// 写入不同自然日的消息
	// Day 1: 2026-09-01
	// Day 2: 2026-09-02
	// ...
	// Day 10: 2026-09-10 (包含 305 条消息触发截断)
	shanghaiTZ := time.FixedZone("CST", 8*3600)
	baseDate, _ := time.ParseInLocation("2006-01-02", "2026-09-01", shanghaiTZ)

	var targetAnchorID int64

	for d := 0; d < 10; d++ {
		curDate := baseDate.AddDate(0, 0, d)
		curDayKey := curDate.Format("2006-01-02")
		msgCount := 5
		if d == 9 {
			msgCount = 310 // 触发 MaxMsgPerDay(300) 截断
		}

		for i := 0; i < msgCount; i++ {
			msgTime := curDate.Add(time.Duration(i) * time.Minute)
			meta := map[string]any{
				"day_key":    curDayKey,
				"created_at": msgTime,
			}
			id, err := mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", fmt.Sprintf("Day %s Msg %d", curDayKey, i), "QUERY_PROGRESS", "", meta)
			if err != nil {
				t.Fatalf("RecordUserMessageWithID failed: %v", err)
			}
			if d == 5 && i == 2 {
				targetAnchorID = id
			}
		}
	}

	// 2. 默认拉取 7 天窗口
	res, err := mgr.GetChatWindow(ctx, ctxmgr.ChatWindowQuery{
		SessionID: sessID,
		Days:      7,
	})
	if err != nil {
		t.Fatalf("GetChatWindow failed: %v", err)
	}

	if res.SessionID != sessID {
		t.Errorf("expected SessionID %s, got %s", sessID, res.SessionID)
	}
	// 总共有 10 天，默认拉 7 天，应返回最近的 7 天
	if len(res.Days) != 7 {
		t.Fatalf("expected 7 days loaded, got %d", len(res.Days))
	}
	// 正序检查：最早日在前，最新日在后
	for i := 1; i < len(res.Days); i++ {
		if res.Days[i-1].DayKey >= res.Days[i].DayKey {
			t.Errorf("expected ascending day order: day[%d]=%s should be before day[%d]=%s", i-1, res.Days[i-1].DayKey, i, res.Days[i].DayKey)
		}
	}
	// 验证还有更早日
	if !res.HasMoreBefore {
		t.Errorf("expected has_more_before to be true")
	}
	oldestInWindow := res.OldestLoadedDay

	// 验证最新的一天（Day 10）发生截断
	latestDay := res.Days[len(res.Days)-1]
	if latestDay.Count != 310 {
		t.Errorf("expected Day 10 count 310, got %d", latestDay.Count)
	}
	if !latestDay.Truncated {
		t.Errorf("expected Day 10 truncated to be true")
	}
	if latestDay.LoadedCount != 300 {
		t.Errorf("expected loaded_count 300, got %d", latestDay.LoadedCount)
	}

	// 3. 触顶加载更早日 (before_day)
	olderRes, err := mgr.GetChatWindow(ctx, ctxmgr.ChatWindowQuery{
		SessionID: sessID,
		Days:      7,
		BeforeDay: oldestInWindow,
	})
	if err != nil {
		t.Fatalf("GetChatWindow before_day failed: %v", err)
	}
	// 剩余 3 天
	if len(olderRes.Days) != 3 {
		t.Errorf("expected 3 older days, got %d", len(olderRes.Days))
	}
	// 此段不应再有更早日
	if olderRes.HasMoreBefore {
		t.Errorf("expected olderRes.HasMoreBefore to be false")
	}

	// 4. ListChatDates 统计测试
	dates, err := mgr.ListChatDates(ctx, sessID)
	if err != nil {
		t.Fatalf("ListChatDates failed: %v", err)
	}
	if len(dates) != 10 {
		t.Fatalf("expected 10 dates, got %d", len(dates))
	}
	// 验证分类统计
	foundDay10 := false
	for _, dt := range dates {
		if dt.Count == 310 {
			foundDay10 = true
			if dt.Categories[ctxmgr.CategoryConsult] != 310 {
				t.Errorf("expected 310 consult messages in day 10 stats, got %d", dt.Categories[ctxmgr.CategoryConsult])
			}
		}
	}
	if !foundDay10 {
		t.Errorf("day 10 with 310 messages not found in date stats")
	}

	// 5. GetChatDayMessages 锚点上下文居中测试
	day5Key := baseDate.AddDate(0, 0, 5).Format("2006-01-02")
	day5Res, err := mgr.GetChatDayMessages(ctx, sessID, day5Key, targetAnchorID, 0, 0, 4)
	if err != nil {
		t.Fatalf("GetChatDayMessages failed: %v", err)
	}
	if day5Res.TotalCount != 5 {
		t.Errorf("expected total 5 messages in Day 5, got %d", day5Res.TotalCount)
	}
	// limit 4 -> returned 4
	if len(day5Res.Messages) != 4 {
		t.Errorf("expected 4 messages returned around anchor, got %d", len(day5Res.Messages))
	}
}

func TestChatWindow_DatabaseIntegration(t *testing.T) {
	dsn := os.Getenv("CONTEXT_DB_DSN")
	if dsn == "" {
		dsn = os.Getenv("TEST_DATABASE_URL")
	}
	if dsn == "" {
		t.Skip("CONTEXT_DB_DSN or TEST_DATABASE_URL must be set (refusing default dev database)")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		t.Skipf("cannot connect to context db (%v), skipping DB integration test", err)
		return
	}
	defer pool.Close()

	if err := pool.Ping(ctx); err != nil {
		t.Skipf("context db ping failed (%v), skipping DB integration test", err)
		return
	}

	mgr := ctxmgr.NewContextManager(pool, nil, ctxmgr.DefaultTimeDecayConfig())
	sessID := fmt.Sprintf("sess_db_window_%d", time.Now().UnixNano())

	// 插入跨越两天的测试消息
	shanghaiTZ := time.FixedZone("CST", 8*3600)
	d1 := time.Now().In(shanghaiTZ).AddDate(0, 0, -1)
	d2 := time.Now().In(shanghaiTZ)
	d1Key := d1.Format("2006-01-02")
	d2Key := d2.Format("2006-01-02")

	id1, err := mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "昨天立项需求沟通", "SUPPLEMENT_INFO", "", map[string]any{
		"day_key":    d1Key,
		"created_at": d1,
	})
	if err != nil {
		t.Fatalf("RecordUserMessageWithID d1 failed: %v", err)
	}

	id2, err := mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "今天补充型检报告", "QUERY_REGULATIONS", "", map[string]any{
		"day_key":    d2Key,
		"created_at": d2,
	})
	if err != nil {
		t.Fatalf("RecordUserMessageWithID d2 failed: %v", err)
	}

	_, err = mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "启动注册申报任务", "TRIGGER_TASK", "", map[string]any{
		"day_key":    d2Key,
		"created_at": d2.Add(time.Minute),
	})
	if err != nil {
		t.Fatalf("RecordUserMessageWithID d2 task failed: %v", err)
	}

	// 1. 测试真实 DB 的 GetChatWindow
	winRes, err := mgr.GetChatWindow(ctx, ctxmgr.ChatWindowQuery{
		SessionID: sessID,
		Days:      7,
	})
	if err != nil {
		t.Fatalf("GetChatWindow failed: %v", err)
	}
	if len(winRes.Days) != 2 {
		t.Fatalf("expected 2 days in window, got %d", len(winRes.Days))
	}
	if winRes.Days[0].DayKey != d1Key || winRes.Days[1].DayKey != d2Key {
		t.Errorf("expected sorted days [%s, %s], got [%s, %s]", d1Key, d2Key, winRes.Days[0].DayKey, winRes.Days[1].DayKey)
	}
	if winRes.Days[0].Messages[0].ID != id1 || winRes.Days[1].Messages[0].ID != id2 {
		t.Errorf("expected message IDs [%d, %d], got [%d, %d]", id1, id2, winRes.Days[0].Messages[0].ID, winRes.Days[1].Messages[0].ID)
	}

	// 2. 测试真实 DB 的 ListChatDates
	dateStats, err := mgr.ListChatDates(ctx, sessID)
	if err != nil {
		t.Fatalf("ListChatDates failed: %v", err)
	}
	if len(dateStats) != 2 {
		t.Fatalf("expected 2 date stats, got %d", len(dateStats))
	}

	// 3. 测试真实 DB 的 GetChatDayMessages
	dayMsgs, err := mgr.GetChatDayMessages(ctx, sessID, d1Key, id1, 0, 0, 10)
	if err != nil {
		t.Fatalf("GetChatDayMessages failed: %v", err)
	}
	if dayMsgs.TotalCount != 1 || len(dayMsgs.Messages) != 1 || dayMsgs.Messages[0].ID != id1 {
		t.Errorf("expected day 1 message ID %d, got %v", id1, dayMsgs)
	}

	// 4. 测试真实 DB 的 ListChatCategories
	catStats, err := mgr.ListChatCategories(ctx, sessID)
	if err != nil {
		t.Fatalf("ListChatCategories failed: %v", err)
	}
	if len(catStats) < 2 {
		t.Fatalf("expected at least 2 categories, got %d", len(catStats))
	}
}

func TestListChatCategories_MemoryMode(t *testing.T) {
	mgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.DefaultTimeDecayConfig())
	ctx := context.Background()
	sessID := "sess_mem_cat_test"

	// 注入不同分类的消息
	_, _ = mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "阶段推进", "ADVANCE_STAGE", "", nil) // task
	_, _ = mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "启动任务", "TRIGGER_TASK", "", nil)   // task
	_, _ = mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "合规咨询", "QUERY_REGULATIONS", "", nil) // consult
	_, _ = mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "立项里程碑", "", "", map[string]any{"milestone": true}) // milestone

	stats, err := mgr.ListChatCategories(ctx, sessID)
	if err != nil {
		t.Fatalf("ListChatCategories failed: %v", err)
	}

	if len(stats) != 3 {
		t.Fatalf("expected 3 distinct categories, got %d", len(stats))
	}

	// 按 count DESC 排序，task 应排在首位 (count=2)
	if stats[0].Category != "task" || stats[0].Count != 2 || stats[0].Label != "任务执行" {
		t.Errorf("expected first category to be task (count=2), got %+v", stats[0])
	}
}