package engine

import (
	"testing"

	"git.kingecg.top/kingecg/gomog/pkg/types"
)

func TestAggregationEngine(t *testing.T) {
	store := NewMemoryStore(nil)
	engine := NewAggregationEngine(store)

	// 创建测试数据
	docs := []types.Document{
		{ID: "1", Data: map[string]interface{}{"category": "A", "amount": 100, "status": "completed"}},
		{ID: "2", Data: map[string]interface{}{"category": "A", "amount": 150, "status": "completed"}},
		{ID: "3", Data: map[string]interface{}{"category": "B", "amount": 200, "status": "pending"}},
		{ID: "4", Data: map[string]interface{}{"category": "B", "amount": 50, "status": "completed"}},
	}

	tests := []struct {
		name     string
		pipeline []types.AggregateStage
		wantLen  int
	}{
		{
			name: "match stage",
			pipeline: []types.AggregateStage{
				{Stage: "$match", Spec: map[string]interface{}{"status": "completed"}},
			},
			wantLen: 3,
		},
		{
			name: "group stage",
			pipeline: []types.AggregateStage{
				{Stage: "$group", Spec: map[string]interface{}{
					"_id":   "$category",
					"total": map[string]interface{}{"$sum": "$amount"},
				}},
			},
			wantLen: 2,
		},
		{
			name: "match and group",
			pipeline: []types.AggregateStage{
				{Stage: "$match", Spec: map[string]interface{}{"status": "completed"}},
				{Stage: "$group", Spec: map[string]interface{}{
					"_id":   "$category",
					"total": map[string]interface{}{"$sum": "$amount"},
				}},
			},
			wantLen: 2,
		},
		{
			name: "sort and limit",
			pipeline: []types.AggregateStage{
				{Stage: "$sort", Spec: map[string]interface{}{"amount": -1}},
				{Stage: "$limit", Spec: 2},
			},
			wantLen: 2,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result, err := engine.ExecutePipeline(docs, tt.pipeline)
			if err != nil {
				t.Errorf("ExecutePipeline() error = %v", err)
				return
			}
			if len(result) != tt.wantLen {
				t.Errorf("ExecutePipeline() result length = %d, want %d", len(result), tt.wantLen)
			}
		})
	}
}

// ExecutePipeline 执行管道（用于测试）
func (e *AggregationEngine) ExecutePipeline(docs []types.Document, pipeline []types.AggregateStage) ([]types.Document, error) {
	result := docs
	for _, stage := range pipeline {
		var err error
		result, err = e.executeStage(stage, result)
		if err != nil {
			return nil, err
		}
	}
	return result, nil
}