package impl

import (
	"context"
	"errors"
	"sort"
	"sync"

	"github.com/pingcap-inc/clinic/pkg/ngm/statements/aggregate"
	"github.com/pingcap-inc/clinic/pkg/ngm/statements/api"
	"github.com/pingcap-inc/clinic/pkg/ngm/statements/filter"
	"github.com/pingcap-inc/clinic/pkg/ngm/statements/group"
	"github.com/pingcap-inc/clinic/pkg/ngm/statements/load"
)

const (
	Concurrency = 8
)

var (
	ErrLoadFailed    = errors.New("load failed")
	ErrExecuteFailed = errors.New("execute failed")
)

func Execute[K comparable](
	ctx context.Context,
	records <-chan load.TimeRangeRecords,
	filter filter.Filter,
	group group.Group[K],
	aggregator aggregate.AggregatorFactory,
) ([]*api.Model, error) {
	total := make(chan map[K]aggregate.Aggregator, Concurrency)

	var wg sync.WaitGroup
	for i := 0; i < Concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			res := make(map[K]aggregate.Aggregator)
			for {
				select {
				case <-ctx.Done():
					return
				case trs, ok := <-records:
					if !ok {
						total <- res
						return
					}

					filter.Filter(trs)
					trs.Traverse(func(r *load.StmtRecord) {
						key := group.Key(r)
						if _, ok := res[key]; !ok {
							res[key] = aggregator()
						}
						res[key].Add(r)
					})
				}
			}
		}()
	}

	go func() {
		wg.Wait()
		close(total)
	}()

	groupAgg := make(map[K]aggregate.Aggregator)

OUT:
	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case part, ok := <-total:
			if !ok {
				break OUT
			}

			for k, v := range part {
				if _, ok := groupAgg[k]; !ok {
					groupAgg[k] = aggregator()
				}
				groupAgg[k].Merge(v)
			}
		}
	}

	res := make([]*api.Model, 0, len(groupAgg))
	for _, v := range groupAgg {
		m := &api.Model{}
		v.Output(m)
		res = append(res, m)
	}

	return res, nil
}

func getStmtAggFactories(fields []string) []aggregate.AggregatorFactory {
	factories := make(map[string]aggregate.AggregatorFactory)

	for _, field := range fields {
		if f, ok := aggregate.StatementAggregatorFactories[field]; ok {
			factories[field] = f
		}
	}

	factoriesSlice := make([]aggregate.AggregatorFactory, 0, len(factories))
	for _, factory := range factories {
		factoriesSlice = append(factoriesSlice, factory)
	}

	return factoriesSlice
}

func allAggFactories() []aggregate.AggregatorFactory {
	factories := make([]aggregate.AggregatorFactory, 0, len(aggregate.StatementAggregatorFactories))
	for _, factory := range aggregate.StatementAggregatorFactories {
		factories = append(factories, factory)
	}

	return factories
}

func sortBySumLatency(statements []*api.Model) {
	sort.Slice(statements, func(i, j int) bool {
		return statements[i].SumLatency > statements[j].SumLatency
	})
}