package main

import (
	"context"
	"fmt"
	"log"
	"time"

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

func main() {
	// 创建内存存储和流式聚合引擎
	store := engine.NewMemoryStore(nil)
	aggEngine := engine.NewStreamAggregationEngine(store)

	// 创建测试数据
	collection := "test_stream"
	testDocs := make(map[string]types.Document)
	for i := 0; i < 1000; i++ {
		testDocs[fmt.Sprintf("doc%d", i)] = types.Document{
			ID: fmt.Sprintf("doc%d", i),
			Data: map[string]interface{}{
				"name":   fmt.Sprintf("User%d", i),
				"age":    20 + (i % 50),
				"score":  float64(50 + (i % 50)),
				"status": map[string]interface{}{"active": i%2 == 0},
			},
		}
	}

	// 插入测试数据
	var docs []types.Document
	for _, doc := range testDocs {
		docs = append(docs, doc)
	}

	if err := store.InsertMany(collection, docs); err != nil {
		log.Printf("Error inserting documents: %v", err)
		return
	}

	// 定义聚合管道
	pipeline := []types.AggregateStage{
		{
			Stage: "$match",
			Spec: map[string]interface{}{
				"age": map[string]interface{}{
					"$gte": 25,
					"$lte": 35,
				},
			},
		},
		{
			Stage: "$project",
			Spec: map[string]interface{}{
				"name":  1,
				"age":   1,
				"score": 1,
			},
		},
		{
			Stage: "$addFields",
			Spec: map[string]interface{}{
				"isHighScorer": map[string]interface{}{
					"$gte": []interface{}{"$score", 80.0},
				},
			},
		},
		{
			Stage: "$sort",
			Spec: map[string]interface{}{
				"score": -1,
			},
		},
		{
			Stage: "$limit",
			Spec:  10,
		},
	}

	// 执行流式聚合
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	opts := engine.StreamAggregationOptions{
		BufferSize: 100, // 每次处理100个文档
	}

	resultChan, errChan := aggEngine.StreamExecute(ctx, collection, pipeline, opts)

	// 处理结果
	var finalResults []types.Document
	for {
		select {
		case batch, ok := <-resultChan:
			if !ok {
				resultChan = nil
				continue
			}
			fmt.Printf("Received batch with %d documents\n", len(batch))
			finalResults = append(finalResults, batch...)
		case err, ok := <-errChan:
			if ok && err != nil {
				log.Printf("Error during stream aggregation: %v", err)
				return
			}
			if !ok {
				// 通道已关闭，结束循环
				goto done
			}
		case <-ctx.Done():
			log.Println("Context cancelled")
			return
		}
	}

done:
	fmt.Printf("Total results: %d\n", len(finalResults))
	for i, doc := range finalResults {
		if i < 5 { // 只打印前5个结果
			fmt.Printf("Result %d: ID=%s, Name=%s, Age=%v, Score=%v\n",
				i+1, doc.ID, doc.Data["name"], doc.Data["age"], doc.Data["score"])
		}
	}
}