// Copyright (c) 2016-2017 Tigera, Inc. All rights reserved.

// Licensed 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 labelindex provides the InheritIndex type, which emits events as the set of
// items (currently WorkloadEndpoints/HostEndpoint) it has been told about start (or stop) matching
// the label selectors (which are extracted from the active policy rules) it has been told about.
//
// # Label inheritance
//
// As the name suggests, the InheritIndex supports the notion of label inheritance.  In our
// data-model:
//
//   - endpoints have their own labels; these take priority over any inherited labels
//   - endpoints also inherit labels from any explicitly-named profiles in their data
//   - profiles have explicit labels
//
// For example, suppose an endpoint had labels
//
//	{"a": "ep-a", "b": "ep-b"}
//
// and it explicitly referenced profile "profile-A", which had these labels:
//
//	{"a": "prof-a", "c": "prof-c", "d": "prof-d"}
//
// then the resulting labels for the endpoint after considering inheritance would be:
//
//	{
//	    "a": "ep-a",    // Explicit endpoint label "wins" over profile labels.
//	    "b": "ep-b",
//	    "c": "prof-c",  // Profile label gets inherited.
//	    "d": "prof-d",
//	}
package labelindex

import (
	"reflect"

	v3 "github.com/projectcalico/api/pkg/apis/projectcalico/v3"
	log "github.com/sirupsen/logrus"

	"github.com/projectcalico/calico/lib/std/uniquelabels"
	"github.com/projectcalico/calico/lib/std/uniquestr"
	"github.com/projectcalico/calico/libcalico-go/lib/backend/api"
	"github.com/projectcalico/calico/libcalico-go/lib/backend/model"
	"github.com/projectcalico/calico/libcalico-go/lib/selector"
	"github.com/projectcalico/calico/libcalico-go/lib/selector/parser"
	"github.com/projectcalico/calico/libcalico-go/lib/set"
)

// FIXME make InheritIndex generic

// itemData holds the data that we know about a particular item (i.e. a workload or host endpoint).
// In particular, it holds it current explicitly-assigned labels and a pointer to the parent data
// for each of its parents.
type itemData struct {
	labels  uniquelabels.Map
	parents []*parentData
}

// GetHandle implements the Labels interface for itemData.  Combines the item's own labels with those
// of its parents on the fly.
func (itemData *itemData) GetHandle(labelName uniquestr.Handle) (handle uniquestr.Handle, present bool) {
	if handle, present = itemData.labels.GetHandle(labelName); present {
		return
	}
	for _, parent := range itemData.parents {
		if handle, present = parent.labels.GetHandle(labelName); present {
			return
		}
	}
	return
}

// parentData holds the data that we know about each parent (i.e. each security profile).  Since,
// profiles consist of multiple resources in our data-model, any of the fields may be nil if we
// have partial information.
type parentData struct {
	id      string
	labels  uniquelabels.Map
	itemIDs set.Typed[any]
}

type MatchCallback func(selId, labelId any)

type InheritIndex struct {
	itemDataByID         map[any]*itemData
	parentDataByParentID map[string]*parentData
	selectorsById        map[any]*selector.Selector

	// Current matches.
	selIdsByLabelId map[any]set.Typed[any]
	labelIdsBySelId map[any]set.Typed[any]

	// Callback functions
	OnMatchStarted MatchCallback
	OnMatchStopped MatchCallback

	dirtyItemIDs set.Typed[any]
}

func NewInheritIndex(onMatchStarted, onMatchStopped MatchCallback) *InheritIndex {
	itemData := map[any]*itemData{}
	inheritIDx := InheritIndex{
		itemDataByID:         itemData,
		parentDataByParentID: map[string]*parentData{},
		selectorsById:        map[any]*selector.Selector{},

		selIdsByLabelId: map[any]set.Typed[any]{},
		labelIdsBySelId: map[any]set.Typed[any]{},

		// Callback functions
		OnMatchStarted: onMatchStarted,
		OnMatchStopped: onMatchStopped,

		dirtyItemIDs: set.New[any](),
	}
	return &inheritIDx
}

// OnUpdate makes LabelInheritanceIndex compatible with the UpdateHandler interface
// allowing it to be used in a calculation graph more easily.
func (l *InheritIndex) OnUpdate(update api.Update) (_ bool) {
	switch key := update.Key.(type) {
	case model.WorkloadEndpointKey:
		if update.Value != nil {
			log.Debugf("Updating InheritIndex with endpoint %v", key)
			endpoint := update.Value.(*model.WorkloadEndpoint)
			profileIDs := endpoint.ProfileIDs
			l.UpdateLabels(key, endpoint.Labels, profileIDs)
		} else {
			log.Debugf("Deleting endpoint %v from InheritIndex", key)
			l.DeleteLabels(key)
		}
	case model.HostEndpointKey:
		if update.Value != nil {
			// Figure out what's changed and update the cache.
			log.Debugf("Updating InheritIndex for host endpoint %v", key)
			endpoint := update.Value.(*model.HostEndpoint)
			profileIDs := endpoint.ProfileIDs
			l.UpdateLabels(key, endpoint.Labels, profileIDs)
		} else {
			log.Debugf("Deleting host endpoint %v from InheritIndex", key)
			l.DeleteLabels(key)
		}
	case model.ResourceKey:
		if key.Kind != v3.KindProfile {
			return
		}
		if update.Value != nil {
			log.Debugf("Updating InheritIndex for profile labels %v", key)
			labels := update.Value.(*v3.Profile).Spec.LabelsToApply
			l.UpdateParentLabels(key.Name, labels)
		} else {
			log.Debugf("Removing profile labels %v from InheritIndex", key)
			l.DeleteParentLabels(key.Name)
		}
	}
	return
}

func (idx *InheritIndex) UpdateSelector(id any, sel *selector.Selector) {
	if sel == nil {
		log.WithField("id", id).Panic("Selector should not be nil")
		panic("Selector should not be nil")
	}
	oldSel := idx.selectorsById[id]
	// Since the selectorRoot struct has cache fields, the easiest way to compare two
	// selectors is to compare their IDs.
	if oldSel.Equal(sel) {
		log.WithField("selID", id).Debug("Skipping unchanged selector")
		return
	}
	log.WithFields(log.Fields{
		"id":       id,
		"selector": sel,
	}).Info("Updating selector")
	idx.scanAllLabels(id, sel)
	idx.selectorsById[id] = sel
}

func (idx *InheritIndex) DeleteSelector(id any) {
	log.Infof("Deleting selector %v", id)
	matchSet := idx.labelIdsBySelId[id]
	if matchSet != nil {
		for labelId := range matchSet.All() {
			// This modifies the set we're iterating over, but that's safe in Go.
			idx.deleteMatch(id, labelId)
		}
	}
	delete(idx.selectorsById, id)
}

func (idx *InheritIndex) UpdateLabels(id any, labels uniquelabels.Map, parentIDs []string) {
	log.Debug("Inherit index updating labels for ", id)
	log.Debug("Num dirty items ", idx.dirtyItemIDs.Len(), " items")

	oldItemData := idx.itemDataByID[id]
	var oldParents []*parentData
	if oldItemData != nil {
		oldParents = oldItemData.parents
		oldLabels := oldItemData.labels
		if reflect.DeepEqual(oldLabels, labels) &&
			reflect.DeepEqual(oldParents, parentIDs) {
			log.Debug("No change to labels or parentIDs, ignoring.")
			return
		}
	}
	newItemData := &itemData{}
	if labels.Len() > 0 {
		newItemData.labels = labels
	}
	if len(parentIDs) > 0 {
		parents := make([]*parentData, len(parentIDs))
		for i, pID := range parentIDs {
			parents[i] = idx.getOrCreateParent(pID)
		}
		newItemData.parents = parents
	}
	idx.itemDataByID[id] = newItemData

	idx.onItemParentsUpdate(id, oldParents, newItemData.parents)

	idx.dirtyItemIDs.Add(id)
	idx.flushUpdates()
	log.Debug("Num ending dirty items ", idx.dirtyItemIDs.Len(), " items")
}

func (idx *InheritIndex) DeleteLabels(id any) {
	log.Debug("Inherit index deleting labels for ", id)
	oldItemData := idx.itemDataByID[id]
	var oldParents []*parentData
	if oldItemData != nil {
		oldParents = oldItemData.parents
	}
	delete(idx.itemDataByID, id)
	idx.onItemParentsUpdate(id, oldParents, nil)
	idx.dirtyItemIDs.Add(id)
	idx.flushUpdates()
}

func (idx *InheritIndex) getOrCreateParent(id string) *parentData {
	parent := idx.parentDataByParentID[id]
	if parent == nil {
		parent = &parentData{
			id: id,
		}
		idx.parentDataByParentID[id] = parent
	}
	return parent
}

func (idx *InheritIndex) discardParentIfEmpty(id string) {
	parent := idx.parentDataByParentID[id]
	if parent == nil {
		return
	}
	if parent.itemIDs == nil && parent.labels.IsNil() {
		delete(idx.parentDataByParentID, id)
	}
}

func (idx *InheritIndex) onItemParentsUpdate(id any, oldParents, newParents []*parentData) {
	log.WithFields(log.Fields{
		"oldParents": oldParents,
		"newParents": newParents,
		"id":         id,
	}).Debug("Updating parents")
	// Calculate the current set of parent IDs so we can skip deletion of parents that are still
	// present.  We need to do this to avoid removing a still-current parent via
	// discardParentIfEmpty().
	currentParentIDs := set.New[any]()
	for _, parentData := range newParents {
		currentParentIDs.Add(parentData.id)
	}

	for _, parent := range oldParents {
		if currentParentIDs.Contains(parent.id) {
			// Make sure we don't delete current parents from the index.
			continue
		}
		parent.itemIDs.Discard(id)
		if parent.itemIDs.Len() == 0 {
			parent.itemIDs = nil
		}
		idx.discardParentIfEmpty(parent.id)
	}

	for _, parent := range newParents {
		if parent.itemIDs == nil {
			parent.itemIDs = set.New[any]()
		}
		parent.itemIDs.Add(id)
	}
}

func (idx *InheritIndex) UpdateParentLabels(parentID string, labels map[string]string) {
	parent := idx.getOrCreateParent(parentID)
	parent.labels = uniquelabels.Make(labels) // FIXME intern further upstream?
	idx.flushChildren(parentID)
}

func (idx *InheritIndex) DeleteParentLabels(parentID string) {
	parent := idx.parentDataByParentID[parentID]
	if parent == nil {
		return
	}
	parent.labels = uniquelabels.Map{}
	idx.discardParentIfEmpty(parentID)
	idx.flushChildren(parentID)
}

func (idx *InheritIndex) flushChildren(parentID string) {
	parentData := idx.parentDataByParentID[parentID]
	if parentData != nil && parentData.itemIDs != nil {
		for itemID := range parentData.itemIDs.All() {
			log.Debug("Marking child ", itemID, " dirty")
			idx.dirtyItemIDs.Add(itemID)
		}
	}
	idx.flushUpdates()
}

func (idx *InheritIndex) flushUpdates() {
	for itemID := range idx.dirtyItemIDs.All() {
		log.Debugf("Flushing %#v", itemID)
		_, ok := idx.itemDataByID[itemID]
		if !ok {
			// Item deleted.
			log.Debugf("Flushing delete of item %v", itemID)
			matchSet := idx.selIdsByLabelId[itemID]
			if matchSet != nil {
				for selId := range matchSet.All() {
					// This modifies the set we're iterating over, but that's safe in Go.
					idx.deleteMatch(selId, itemID)
				}
			}
		} else {
			// Item updated/created, re-evaluate labels.
			log.Debugf("Flushing update of item %v", itemID)
			idx.scanAllSelectors(itemID)
		}
		idx.dirtyItemIDs.Discard(itemID)
	}
}

func (idx *InheritIndex) scanAllLabels(selId any, sel *selector.Selector) {
	log.Debugf("Scanning all (%v) labels against selector %v",
		len(idx.itemDataByID), selId)
	for labelId, labels := range idx.itemDataByID {
		idx.updateMatches(selId, sel, labelId, labels)
	}
}

func (idx *InheritIndex) scanAllSelectors(labelId any) {
	log.Debugf("Scanning all (%v) selectors against labels %v",
		len(idx.selectorsById), labelId)
	labels := idx.itemDataByID[labelId]
	for selId, sel := range idx.selectorsById {
		idx.updateMatches(selId, sel, labelId, labels)
	}
}

func (idx *InheritIndex) updateMatches(
	selId any,
	sel *selector.Selector,
	labelId any,
	labels parser.Labels,
) {
	nowMatches := sel.EvaluateLabels(labels)
	if nowMatches {
		idx.storeMatch(selId, labelId)
	} else {
		idx.deleteMatch(selId, labelId)
	}
}

func (idx *InheritIndex) storeMatch(selId, labelId any) {
	labelIds := idx.labelIdsBySelId[selId]
	if labelIds == nil {
		labelIds = set.New[any]()
		idx.labelIdsBySelId[selId] = labelIds
	}
	previouslyMatched := labelIds.Contains(labelId)
	if !previouslyMatched {
		log.Debugf("Selector %v now matches labels %v", selId, labelId)
		labelIds.Add(labelId)

		selIDs, ok := idx.selIdsByLabelId[labelId]
		if !ok {
			selIDs = set.New[any]()
			idx.selIdsByLabelId[labelId] = selIDs
		}
		selIDs.Add(selId)

		idx.OnMatchStarted(selId, labelId)
	}
}

func (idx *InheritIndex) deleteMatch(selId, labelId any) {
	labelIds := idx.labelIdsBySelId[selId]
	if labelIds == nil {
		return
	}
	previouslyMatched := labelIds.Contains(labelId)
	if previouslyMatched {
		log.Debugf("Selector %v no longer matches labels %v",
			selId, labelId)

		labelIds.Discard(labelId)
		if labelIds.Len() == 0 {
			delete(idx.labelIdsBySelId, selId)
		}

		idx.selIdsByLabelId[labelId].Discard(selId)
		if idx.selIdsByLabelId[labelId].Len() == 0 {
			delete(idx.selIdsByLabelId, labelId)
		}

		idx.OnMatchStopped(selId, labelId)
	}
}
