// 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 trace

import (
	"github.com/apache/skywalking-banyandb/api/common"
	"github.com/apache/skywalking-banyandb/banyand/internal/sidx"
	"github.com/apache/skywalking-banyandb/banyand/internal/storage"
	"github.com/apache/skywalking-banyandb/banyand/internal/wqueue"
	"github.com/apache/skywalking-banyandb/pkg/encoding"
	"github.com/apache/skywalking-banyandb/pkg/index"
	pbv1 "github.com/apache/skywalking-banyandb/pkg/pb/v1"
	"github.com/apache/skywalking-banyandb/pkg/pool"
	"github.com/apache/skywalking-banyandb/pkg/timestamp"
)

type tagValue struct {
	tag       string
	value     []byte
	valueArr  [][]byte
	valueType pbv1.ValueType
}

func (t *tagValue) reset() {
	t.tag = ""
	t.value = nil
	t.valueArr = nil
}

func (t *tagValue) marshal() []byte {
	if t.valueArr != nil {
		var dst []byte
		for i := range t.valueArr {
			if t.valueType == pbv1.ValueTypeInt64Arr {
				dst = append(dst, t.valueArr[i]...)
				continue
			}
			dst = encoding.MarshalVarArray(dst, t.valueArr[i])
		}
		return dst
	}
	return t.value
}

func generateTagValue() *tagValue {
	v := tagValuePool.Get()
	if v == nil {
		return &tagValue{}
	}
	return v
}

func releaseTagValue(v *tagValue) {
	v.reset()
	tagValuePool.Put(v)
}

var tagValuePool = pool.Register[*tagValue]("trace-tagValue")

type traces struct {
	traceIDs   []string
	timestamps []int64
	tags       [][]*tagValue
	spans      [][]byte
	spanIDs    []string
}

func (t *traces) reset() {
	t.traceIDs = t.traceIDs[:0]
	t.timestamps = t.timestamps[:0]
	for i := range t.tags {
		for j := range t.tags[i] {
			releaseTagValue(t.tags[i][j])
		}
	}
	t.tags = t.tags[:0]
	t.spans = t.spans[:0]
	t.spanIDs = t.spanIDs[:0]
}

func (t *traces) Len() int {
	return len(t.traceIDs)
}

func (t *traces) Less(i, j int) bool {
	return t.traceIDs[i] < t.traceIDs[j]
}

func (t *traces) Swap(i, j int) {
	t.traceIDs[i], t.traceIDs[j] = t.traceIDs[j], t.traceIDs[i]
	t.timestamps[i], t.timestamps[j] = t.timestamps[j], t.timestamps[i]
	t.tags[i], t.tags[j] = t.tags[j], t.tags[i]
	t.spans[i], t.spans[j] = t.spans[j], t.spans[i]
	t.spanIDs[i], t.spanIDs[j] = t.spanIDs[j], t.spanIDs[i]
}

func generateTraces() *traces {
	v := tracesPool.Get()
	if v == nil {
		return &traces{}
	}
	return v
}

func releaseTraces(t *traces) {
	t.reset()
	tracesPool.Put(t)
}

var tracesPool = pool.Register[*traces]("trace-traces")

type seriesDoc struct {
	docIDsAdded map[uint64]struct{}
	docs        index.Documents
}

type tracesInTable struct {
	segment     storage.Segment[*tsTable, option]
	tsTable     *tsTable
	traces      *traces
	sidxReqsMap map[string][]sidx.WriteRequest
	timeRange   timestamp.TimeRange
	seriesDocs  seriesDoc
	shardID     common.ShardID
}

type tracesInGroup struct {
	tsdb     storage.TSDB[*tsTable, option]
	tables   []*tracesInTable
	segments []storage.Segment[*tsTable, option]
	latestTS int64
}

type tracesInQueue struct {
	name   string
	queue  *wqueue.Queue[*tsTable, option]
	tables []*tracesInTable
}
