package main

import (
	"fmt"
	"reflect"
)

// Person 结构体
type Person struct {
	Name    string
	Country string
	City    string
}

// GroupByBuilder 用于构建多级分组
type GroupByBuilder[T any] struct {
	data     []T
	groups   map[string]interface{}
	keyFuncs []func(T) (string, string)
}

// NewGroupByBuilder 创建一个新的 GroupByBuilder 实例
func NewGroupByBuilder[T any](data []T) *GroupByBuilder[T] {
	return &GroupByBuilder[T]{
		data: data,
	}
}

// Group 添加一个分组键函数
func (b *GroupByBuilder[T]) Group(fieldName string) *GroupByBuilder[T] {
	keyFunc := func(item T) (string, string) {
		v := reflect.ValueOf(item)
		f := v.FieldByName(fieldName)
		if !f.IsValid() {
			panic(fmt.Sprintf("Field %s does not exist in the provided type", fieldName))
		}
		return fieldName, f.String()
	}
	b.keyFuncs = append(b.keyFuncs, keyFunc)
	return b
}

// Build 构建分组
func (b *GroupByBuilder[T]) Build() *GroupByResult[T] {
	if len(b.keyFuncs) == 0 {
		return &GroupByResult[T]{groups: map[string]interface{}{}}
	}

	groups := make(map[string]interface{})
	for _, item := range b.data {
		currentLevel := groups
		for _, keyFunc := range b.keyFuncs {
			_, fieldValue := keyFunc(item)
			if currentLevel[fieldValue] == nil {
				currentLevel[fieldValue] = make(map[string]interface{})
			}
			nextLevel, ok := currentLevel[fieldValue].(map[string]interface{})
			if !ok {
				panic(fmt.Sprintf("Expected map for key %s, got %T", fieldValue, currentLevel[fieldValue]))
			}
			currentLevel = nextLevel
		}
		// 最后一层添加数据
		if currentLevel["items"] == nil {
			currentLevel["items"] = make([]T, 0)
		}
		group, ok := currentLevel["items"].([]T)
		if !ok {
			panic(fmt.Sprintf("Expected slice for key items, got %T", currentLevel["items"]))
		}
		group = append(group, item)
		currentLevel["items"] = group
	}

	return &GroupByResult[T]{groups: groups}
}

// GroupByResult 存储分组结果
type GroupByResult[T any] struct {
	groups map[string]interface{}
}

// GetAll 获取每个分组的所有元素
func (r *GroupByResult[T]) GetAll() map[string]interface{} {
	return r.groups
}

// 示例数据
var people = []Person{
	{Name: "Alice", Country: "USA", City: "New York"},
	{Name: "Bob", Country: "USA", City: "Los Angeles"},
	{Name: "Charlie", Country: "Canada", City: "Toronto"},
	{Name: "David", Country: "Canada", City: "Vancouver"},
	{Name: "Eve", Country: "USA", City: "New York"},
	{Name: "Frank", Country: "Canada", City: "Toronto"},
}

func main() {
	// 创建 GroupByBuilder 实例
	builder := NewGroupByBuilder(people)

	// 按 Country 和 City 分组
	grouped := builder.Group("Country").Build()

	// 打印每个分组的所有元素
	allElements := grouped.GetAll()
	fmt.Println("\nAll elements in each group:")
	printGroups(allElements, "")
}

// printGroups 递归打印分组
func printGroups(groups map[string]interface{}, prefix string) {
	for k, v := range groups {
		switch v := v.(type) {
		case []Person:
			fmt.Printf("%s%s: ", prefix, k)
			for _, person := range v {
				fmt.Printf("%s, ", person.Name)
			}
			fmt.Println()
		case map[string]interface{}:
			newPrefix := prefix + k + " -> "
			printGroups(v, newPrefix)
		default:
			panic(fmt.Sprintf("Unexpected type: %T", v))
		}
	}
}