// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The 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 pipeline

import (
	"go.uber.org/zap"

	"github.com/milvus-io/milvus/internal/proto/datapb"
	"github.com/milvus-io/milvus/internal/querynodev2/delegator"
	base "github.com/milvus-io/milvus/internal/util/pipeline"
	"github.com/milvus-io/milvus/pkg/log"
	"github.com/milvus-io/milvus/pkg/metrics"
	"github.com/milvus-io/milvus/pkg/mq/msgdispatcher"
	"github.com/milvus-io/milvus/pkg/util/paramtable"
	"github.com/milvus-io/milvus/pkg/util/typeutil"
)

// pipeline used for querynode
type Pipeline interface {
	base.StreamPipeline
	ExcludedSegments(segInfos ...*datapb.SegmentInfo)
}

type pipeline struct {
	base.StreamPipeline

	excludedSegments *typeutil.ConcurrentMap[int64, *datapb.SegmentInfo]
	collectionID     UniqueID
}

func (p *pipeline) ExcludedSegments(segInfos ...*datapb.SegmentInfo) {
	for _, segInfo := range segInfos {
		log.Debug("pipeline add exclude info",
			zap.Int64("segmentID", segInfo.GetID()),
			zap.Uint64("ts", segInfo.GetDmlPosition().GetTimestamp()),
		)
		p.excludedSegments.Insert(segInfo.GetID(), segInfo)
	}
}

func (p *pipeline) Close() {
	p.StreamPipeline.Close()
	metrics.CleanupQueryNodeCollectionMetrics(paramtable.GetNodeID(), p.collectionID)
}

func NewPipeLine(
	collectionID UniqueID,
	channel string,
	manager *DataManager,
	tSafeManager TSafeManager,
	dispatcher msgdispatcher.Client,
	delegator delegator.ShardDelegator,
) (Pipeline, error) {
	pipelineQueueLength := paramtable.Get().QueryNodeCfg.FlowGraphMaxQueueLength.GetAsInt32()
	excludedSegments := typeutil.NewConcurrentMap[int64, *datapb.SegmentInfo]()

	p := &pipeline{
		collectionID:     collectionID,
		excludedSegments: excludedSegments,
		StreamPipeline:   base.NewPipelineWithStream(dispatcher, nodeCtxTtInterval, enableTtChecker, channel),
	}

	filterNode := newFilterNode(collectionID, channel, manager, excludedSegments, pipelineQueueLength)
	insertNode := newInsertNode(collectionID, channel, manager, delegator, pipelineQueueLength)
	deleteNode := newDeleteNode(collectionID, channel, manager, tSafeManager, delegator, pipelineQueueLength)
	p.Add(filterNode, insertNode, deleteNode)
	return p, nil
}