// Licensed to Apache Software Foundation (ASF) under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Apache Software Foundation (ASF) licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package dquery

import (
	"context"
	"errors"
	"fmt"
	"time"

	"go.uber.org/multierr"

	"github.com/apache/skywalking-banyandb/api/common"
	"github.com/apache/skywalking-banyandb/api/data"
	commonv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/common/v1"
	databasev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/database/v1"
	measurev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/measure/v1"
	modelv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/model/v1"
	"github.com/apache/skywalking-banyandb/banyand/measure"
	"github.com/apache/skywalking-banyandb/pkg/bus"
	"github.com/apache/skywalking-banyandb/pkg/convert"
	"github.com/apache/skywalking-banyandb/pkg/flow/streaming"
	"github.com/apache/skywalking-banyandb/pkg/iter/sort"
	"github.com/apache/skywalking-banyandb/pkg/logger"
	pbv1 "github.com/apache/skywalking-banyandb/pkg/pb/v1"
	pkgquery "github.com/apache/skywalking-banyandb/pkg/query"
)

const defaultTopNQueryTimeout = 10 * time.Second

type topNQueryProcessor struct {
	measureService measure.Service
	broadcaster    bus.Broadcaster
	*queryService
	*bus.UnImplementedHealthyListener
}

func (t *topNQueryProcessor) getTopNFieldType(ctx context.Context, group, topNName string) (databasev1.FieldType, error) {
	topNAgg, err := t.metaService.TopNAggregationRegistry().GetTopNAggregation(ctx, &commonv1.Metadata{
		Group: group,
		Name:  topNName,
	})
	if err != nil {
		return databasev1.FieldType_FIELD_TYPE_UNSPECIFIED, fmt.Errorf("failed to get topN aggregation %s/%s: %w", group, topNName, err)
	}
	sourceMeasure := topNAgg.GetSourceMeasure()
	measureSchema, err := t.metaService.MeasureRegistry().GetMeasure(ctx, sourceMeasure)
	if err != nil {
		return databasev1.FieldType_FIELD_TYPE_UNSPECIFIED, fmt.Errorf("failed to get source measure %s: %w", sourceMeasure.GetName(), err)
	}
	fieldName := topNAgg.GetFieldName()
	for _, field := range measureSchema.GetFields() {
		if field.GetName() == fieldName {
			ft := field.GetFieldType()
			if ft != databasev1.FieldType_FIELD_TYPE_INT && ft != databasev1.FieldType_FIELD_TYPE_FLOAT {
				return databasev1.FieldType_FIELD_TYPE_UNSPECIFIED, fmt.Errorf("unsupported field type %s for field %s in topN aggregation", ft.String(), fieldName)
			}
			return ft, nil
		}
	}
	return databasev1.FieldType_FIELD_TYPE_UNSPECIFIED, fmt.Errorf("field %s not found in measure %s", fieldName, sourceMeasure.GetName())
}

func (t *topNQueryProcessor) Rev(ctx context.Context, message bus.Message) (resp bus.Message) {
	request, ok := message.Data().(*measurev1.TopNRequest)
	if !ok {
		t.log.Warn().Msg("invalid event data type")
		return
	}
	n := time.Now()
	now := bus.MessageID(request.TimeRange.Begin.Nanos)
	if err := t.validateRequest(request); err != nil {
		resp = bus.NewMessage(now, common.NewError("%s", err.Error()))
		return
	}
	if e := t.log.Debug(); e.Enabled() {
		e.RawJSON("req", logger.Proto(request)).Msg("received a topN query event")
	}
	nodeSelectors := make(map[string][]string)
	for _, g := range request.Groups {
		if gs, ok := t.measureService.LoadGroup(g); ok {
			if ns, exist := t.parseNodeSelector(request.Stages, gs.GetSchema().ResourceOpts); exist {
				nodeSelectors[g] = ns
			} else if len(gs.GetSchema().ResourceOpts.Stages) > 0 {
				t.log.Error().Strs("req_stages", request.Stages).Strs("default_stages", gs.GetSchema().GetResourceOpts().GetDefaultStages()).Msg("no stage found")
				resp = bus.NewMessage(now, common.NewError("no stage found in request or default stages in resource opts"))
				return
			}
		} else {
			t.log.Error().Str("group", g).Msg("failed to load group")
			resp = bus.NewMessage(now, common.NewError("failed to load group %s", g))
			return
		}
	}
	if len(request.Stages) > 0 && len(nodeSelectors) == 0 {
		t.log.Error().RawJSON("req", logger.Proto(request)).Msg("no stage found")
		resp = bus.NewMessage(now, common.NewError("no stage found"))
		return
	}
	var span *pkgquery.Span
	if request.Trace {
		var tracer *pkgquery.Tracer
		tracer, ctx = pkgquery.NewTracer(ctx, n.Format(time.RFC3339Nano))
		span, _ = tracer.StartSpan(ctx, "distributed-client")
		span.Tag("request", convert.BytesToString(logger.Proto(request)))
		span.Tagf("nodeSelectors", "%v", nodeSelectors)
		defer func() {
			data := resp.Data()
			switch d := data.(type) {
			case *measurev1.TopNResponse:
				d.Trace = tracer.ToProto()
			case *common.Error:
				span.Error(errors.New(d.Error()))
				resp = bus.NewMessage(now, &measurev1.TopNResponse{Trace: tracer.ToProto()})
			default:
				panic("unexpected data type")
			}
			span.Stop()
		}()
	}
	fieldType, fieldTypeErr := t.getTopNFieldType(ctx, request.Groups[0], request.GetName())
	if fieldTypeErr != nil {
		resp = bus.NewMessage(now, common.NewError("failed to determine field type for topN query %s: %v", request.GetName(), fieldTypeErr))
		return
	}
	agg := request.Agg
	request.Agg = modelv1.AggregationFunction_AGGREGATION_FUNCTION_UNSPECIFIED
	ff, err := t.broadcaster.Broadcast(defaultTopNQueryTimeout, data.TopicTopNQuery, bus.NewMessageWithNodeSelectors(now, nodeSelectors, request.TimeRange, request))
	if err != nil {
		resp = bus.NewMessage(now, common.NewError("execute the query %s: %v", request.GetName(), err))
		return
	}
	var allErr error
	var lists []*measurev1.TopNList
	var responseCount int
	if fieldType == databasev1.FieldType_FIELD_TYPE_FLOAT {
		lists, responseCount, allErr = processTopNResponse(ff, request.GetTopN(), agg, request.GetFieldValueSort(),
			measure.CreateTopNPostProcessorFloat, measure.FieldValueToFloat)
	} else {
		lists, responseCount, allErr = processTopNResponse(ff, request.GetTopN(), agg, request.GetFieldValueSort(),
			measure.CreateTopNPostProcessorInt, measure.FieldValueToInt)
	}
	if span != nil {
		span.Tagf("response_count", "%d", responseCount)
	}
	if allErr != nil {
		resp = bus.NewMessage(now, common.NewError("execute the query %s: %v", request.GetName(), allErr))
		return
	}
	if len(lists) == 0 {
		resp = bus.NewMessage(now, &measurev1.TopNResponse{})
		return
	}

	if span != nil {
		span.Tagf("list_count", "%d", len(lists))
	}
	resp = bus.NewMessage(now, &measurev1.TopNResponse{
		Lists: lists,
	})
	if !request.Trace && t.slowQuery > 0 {
		latency := time.Since(n)
		if latency > t.slowQuery {
			t.log.Warn().Dur("latency", latency).RawJSON("req", logger.Proto(request)).Int("resp_count", len(lists)).Msg("top_n slow query")
		}
	}
	return
}

var _ sort.Comparable = (*comparableTopNItem)(nil)

type comparableTopNItem struct {
	*measurev1.TopNList_Item
}

func (c *comparableTopNItem) SortedField() []byte {
	if c.Value.GetFloat() != nil {
		return convert.Float64ToOrderedBytes(c.Value.GetFloat().GetValue())
	}
	return convert.Int64ToBytes(c.Value.GetInt().Value)
}

var _ sort.Iterator[*comparableTopNItem] = (*sortedTopNList)(nil)

type sortedTopNList struct {
	*measurev1.TopNList
	index int
}

func (*sortedTopNList) Close() error {
	return nil
}

func (s *sortedTopNList) Next() bool {
	if s.index >= len(s.Items) {
		return false
	}
	s.index++
	return s.index < len(s.Items)
}

func (s *sortedTopNList) Val() *comparableTopNItem {
	return &comparableTopNItem{s.Items[s.index-1]}
}

func processTopNResponse[K streaming.TopSortKey](
	ff []bus.Future,
	topN int32,
	agg modelv1.AggregationFunction,
	sort modelv1.Sort,
	createAggregator func(int32, modelv1.AggregationFunction, modelv1.Sort) measure.PostProcessor[K],
	convertValue func(*modelv1.FieldValue) K,
) ([]*measurev1.TopNList, int, error) {
	aggregator := createAggregator(topN, agg, sort)
	var tags []string
	var allErr error
	var responseCount int
	for _, f := range ff {
		if m, getErr := f.Get(); getErr != nil {
			allErr = multierr.Append(allErr, getErr)
		} else {
			d := m.Data()
			if d == nil {
				continue
			}
			responseCount++
			topNResp := d.(*measurev1.TopNResponse)
			for _, l := range topNResp.Lists {
				for _, tn := range l.Items {
					if tags == nil {
						tags = make([]string, 0, len(tn.Entity))
						for _, e := range tn.Entity {
							tags = append(tags, e.Key)
						}
					}
					entityValues := make(pbv1.EntityValues, 0, len(tn.Entity))
					for _, e := range tn.Entity {
						entityValues = append(entityValues, e.Value)
					}
					aggregator.Put(entityValues, convertValue(tn.Value), uint64(tn.Timestamp.AsTime().UnixMilli()), tn.Version)
				}
			}
		}
	}
	if tags == nil {
		return nil, responseCount, allErr
	}
	lists, err := aggregator.Val(tags)
	if err != nil {
		return nil, responseCount, multierr.Append(allErr, err)
	}
	return lists, responseCount, allErr
}

func (t *topNQueryProcessor) validateRequest(request *measurev1.TopNRequest) error {
	if request.GetFieldValueSort() == modelv1.Sort_SORT_UNSPECIFIED {
		return errors.New("unspecified requested sort direction")
	}
	if request.GetAgg() == modelv1.AggregationFunction_AGGREGATION_FUNCTION_UNSPECIFIED {
		return errors.New("unspecified requested aggregation function")
	}
	if len(request.GetGroups()) == 0 {
		return errors.New("at least one group is required")
	}
	if request.GetTopN() <= 0 {
		return errors.New("topN must be positive")
	}
	return nil
}
