package watch_graph //nolint:staticcheck

import (
	"context"
	"fmt"
	"log/slog"
	"sync"
	"unique"

	"github.com/google/cel-go/cel"
	"github.com/google/cel-go/common/types"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/k8stool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/discovery"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/jsonpath"
	"k8s.io/client-go/util/workqueue"
)

type ArcType byte // enough bits to represent each arc type as an individual bit.

const (
	UnknownArcType        ArcType = 0 // must be zero
	OwnerReferenceArcType ArcType = 1 << (iota - 1)
	ReferenceArcType
	TransitiveReferenceArcType
)

func (at ArcType) String() string {
	switch at { //nolint:exhaustive
	case OwnerReferenceArcType:
		return "or"
	case ReferenceArcType:
		return "r"
	case TransitiveReferenceArcType:
		return "t"
	default:
		return "unknown"
	}
}

func ParseArcTypeStr(at string) ArcType {
	switch at {
	case "or":
		return OwnerReferenceArcType
	case "r":
		return ReferenceArcType
	case "t":
		return TransitiveReferenceArcType
	default:
		return UnknownArcType
	}
}

// ArcTypeSet is a set of arc types.
type ArcTypeSet ArcType

func (s *ArcTypeSet) Add(v ArcType) {
	*s |= ArcTypeSet(v)
}

func (s *ArcTypeSet) Remove(v ArcType) {
	*s &= ^ArcTypeSet(v)
}

func (s *ArcTypeSet) Contains(v ArcType) bool {
	return (*s & ArcTypeSet(v)) != 0
}

func (s *ArcTypeSet) IsEmpty() bool {
	return *s == 0
}

type gkInfo struct {
	gvr        unique.Handle[schema.GroupVersionResource]
	namespaced bool
}

type gvrInfo struct {
	object     QueryIncludeObject
	namespaced bool
}

type gvrInfos struct {
	infos map[unique.Handle[schema.GroupVersionResource]]gvrInfo

	// Our little simple replacement for the rest mapper for selected resources.
	gk2info map[schema.GroupKind]gkInfo

	// All group-kinds that we got from discovery.
	// Includes non-watchable resources.
	// Does not include sub-resources.
	gks sets.Set[schema.GroupKind]

	// Is discovery info partial or complete.
	// Partial info may not contain all GVRs because e.g. one of the configured aggregated API servers is offline
	// and hence discovery info for the types it exposes is unavailable at the moment.
	// See https://kubernetes.io/docs/concepts/extend-kubernetes/api-extension/apiserver-aggregation/.
	isPartial bool
}

// QueryIncludeObject holds config for object filtering.
type QueryIncludeObject struct {
	LabelSelector            string
	FieldSelector            string
	ObjectSelectorExpression cel.Program        // can be nil
	JSONPath                 *jsonpath.JSONPath // can be nil
}

type QueryInclude struct {
	ResourceSelectorExpression cel.Program // never nil

	Object QueryIncludeObject
}

type QueryExclude struct {
	ResourceSelectorExpression cel.Program // never nil
}

type Namespaces struct {
	Names                    sets.Set[string]
	LabelSelector            string
	FieldSelector            string
	ObjectSelectorExpression cel.Program // can be nil
}

// All checks if all namespaces should be watched.
func (n *Namespaces) All() bool {
	return n.Names.Len() == 0 && n.LabelSelector == "" && n.ObjectSelectorExpression == nil
}

func (n *Namespaces) isFixedList() bool {
	return n.Names.Len() > 0
}

type selectedGVR struct {
	inf  cache.SharedIndexInformer
	stop context.CancelFunc
	wait func()
}

type nsStatus byte

const (
	nsDeleted nsStatus = iota // must be zero
	nsSelected
	nsFilteredOut
)

type nsWithStatus struct {
	name   string
	status nsStatus
}

type VertexID struct {
	GVR       unique.Handle[schema.GroupVersionResource]
	Namespace string
	Name      string
}

type VertexData struct {
	Object   map[string]any
	JSONPath *jsonpath.JSONPath // nil if no JSON Path processing is necessary
}

type ArcAttrs struct {
	Controller              bool `json:"c,omitempty"`
	BlockOwnerDeletion      bool `json:"b,omitempty"`
	DestinationDoesNotExist bool `json:"e,omitempty"`
}

func (a ArcAttrs) IsZero() bool {
	var zero ArcAttrs
	return a == zero
}

type Opts struct {
	Log              *slog.Logger
	Queries          []any // Elements are QueryInclude or QueryExclude
	Namespaces       Namespaces
	DiscoClient      discovery.AggregatedDiscoveryInterface
	Client           dynamic.Interface
	OnWarning        func(Warning)
	ObjectToEvalVars func(*unstructured.Unstructured, schema.GroupVersionResource) map[string]any
	Graph            *ObjectGraph[VertexID, ArcType, VertexData, ArcAttrs]
}

type WatchGraph struct {
	opts      Opts
	allNS     bool
	workqueue workqueue.TypedRateLimitingInterface[VertexID]
	wg        wait.Group

	// Fields below are mutable but are only accessed from the Run goroutine.

	ns           map[string]nsStatus
	selectedGVRs map[unique.Handle[schema.GroupVersionResource]]map[string]selectedGVR // GVR -> namespace -> data
	gvrInfos     gvrInfos
}

func New(opts Opts) *WatchGraph {
	return &WatchGraph{
		opts:  opts,
		allNS: opts.Namespaces.All(),
		workqueue: workqueue.NewTypedRateLimitingQueueWithConfig(
			workqueue.DefaultTypedControllerRateLimiter[VertexID](),
			workqueue.TypedRateLimitingQueueConfig[VertexID]{
				Name: "WatchGraph",
			},
		),
		ns:           map[string]nsStatus{},
		selectedGVRs: map[unique.Handle[schema.GroupVersionResource]]map[string]selectedGVR{},
	}
}

func (g *WatchGraph) Run(ctx context.Context) *Error {
	defer g.wg.Wait()
	defer g.workqueue.ShutDown()
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	done := ctx.Done()
	var nsStatusCh chan nsWithStatus
	errCh := make(chan Error) // fatal errors
	gvrCh := make(chan gvrInfos)
	onError := func(err Error) {
		select {
		case <-done:
		case errCh <- err:
		}
	}
	if !g.allNS {
		nsStatusCh = make(chan nsWithStatus)
		wgn := wgNamespaces{
			namespaces: g.opts.Namespaces,
			client:     g.opts.Client,
			setStatus: func(ns nsWithStatus) {
				select {
				case <-done:
				case nsStatusCh <- ns:
				}
			},
			objectToEvalVars: g.opts.ObjectToEvalVars,
			onError:          onError,
			onWarn:           g.opts.OnWarning,
		}
		g.wg.StartWithContext(ctx, wgn.Run)
	}

	g.wg.Start(func() {
		d := wgDiscovery{
			log:         g.opts.Log,
			queries:     g.opts.Queries,
			discoClient: g.opts.DiscoClient,
			onDiscovery: func(infos gvrInfos) {
				select {
				case <-done:
				case gvrCh <- infos:
				}
			},
			onWarn: g.opts.OnWarning,
		}
		err := d.Run(ctx)
		if err != nil {
			onError(*err)
		}
	})

	wqp := k8stool.NewWorkqueuePuller(g.workqueue)
	initialDiscoveryDone := false
	for {
		select {
		case <-done:
			return nil
		case ns := <-nsStatusCh:
			g.handleNamespaceStatus(ctx, ns)
		case gvr := <-gvrCh:
			g.handleGVRUpdate(ctx, gvr)
			if !initialDiscoveryDone {
				initialDiscoveryDone = true
				g.wg.StartWithContext(ctx, wqp.Run) // start pulling work after we've got discovery info the first time
			}
		case key := <-wqp.Work():
			err := g.processWorkItem(ctx, key)
			if err != nil {
				return err
			}
		case err := <-errCh:
			return &err
		}
	}
}

func (g *WatchGraph) id2object(vid VertexID) (*unstructured.Unstructured, *Error) {
	nsToLookup := metav1.NamespaceAll
	if vid.Namespace != "" && !g.allNS { // it's a namespaced object and we are not watching all namespaces
		nsToLookup = vid.Namespace
	}
	s, ok := g.selectedGVRs[vid.GVR][nsToLookup]
	if !ok {
		return nil, nil
	}
	// Informers use cache.DeletionHandlingMetaNamespaceKeyFunc which uses cache.ObjectName{}.String() to construct
	// key. So we do the same here.
	on := cache.ObjectName{
		Namespace: vid.Namespace,
		Name:      vid.Name,
	}
	key := on.String()
	obj, exists, err := s.inf.GetStore().GetByKey(key)
	if err != nil {
		return nil, &Error{
			Message: fmt.Sprintf("error getting object by key %s from informer: %v", key, err),
			Code:    InternalError,
		}
	}
	if !exists {
		return nil, nil
	}
	return obj.(*unstructured.Unstructured), nil
}

func (g *WatchGraph) isFilteredOutNamespace(ns string) bool {
	if g.opts.Namespaces.isFixedList() && !g.opts.Namespaces.Names.Has(ns) { // Not in the select list
		return true
	}
	return g.ns[ns] == nsFilteredOut
}

func (g *WatchGraph) handleGVRUpdate(ctx context.Context, infos gvrInfos) {
	toStop := g.selectedGVRs

	// Start/stop informers as necessary.
	newSelectedGVRs := make(map[unique.Handle[schema.GroupVersionResource]]map[string]selectedGVR, len(infos.infos))
	for gvr, info := range infos.infos {
		existingData, ok := toStop[gvr]
		if ok { // no changes, keep everything as-is, just move
			newSelectedGVRs[gvr] = existingData
			delete(toStop, gvr) // remove from current set
		} else { // new GVR! Need to start informers for all/selected namespaces.
			ns2inf := make(map[string]selectedGVR)
			if !info.namespaced || g.allNS {
				ns2inf[metav1.NamespaceAll] = g.startInformerForNamespacedGVR(ctx, gvr, info, metav1.NamespaceAll)
			} else {
				for ns, status := range g.ns {
					if status != nsSelected {
						continue
					}
					ns2inf[ns] = g.startInformerForNamespacedGVR(ctx, gvr, info, ns)
				}
			}
			newSelectedGVRs[gvr] = ns2inf
		}
	}
	g.gvrInfos = infos
	g.selectedGVRs = newSelectedGVRs
	var waits []func()
	for _, ns2data := range toStop {
		for _, data := range ns2data {
			data.stop()
			waits = append(waits, data.wait)
		}
	}
	for _, w := range waits {
		w()
	}
}

func (g *WatchGraph) startInformerForNamespacedGVR(ctx context.Context, gvr unique.Handle[schema.GroupVersionResource], info gvrInfo, ns string) selectedGVR {
	g.opts.Log.Debug("Starting informer", logz.GVR(gvr.Value()), logz.NamespacedName(ns))
	infCtx, infCancel := context.WithCancel(ctx)
	nsInf := newInformer(g.opts.Client, gvr.Value(), ns, info.object.LabelSelector, info.object.FieldSelector)
	_, err := nsInf.AddEventHandler(&resourceEventHandler{
		workqueue: g.workqueue,
		gvr:       gvr,
	})
	if err != nil { // cannot happen
		panic(err)
	}
	//nsInf.HasSynced() // TODO handle stuck sync
	var waitWg sync.WaitGroup
	waitWg.Add(1)
	g.wg.Start(func() {
		defer waitWg.Done()
		nsInf.RunWithContext(infCtx)
	})
	return selectedGVR{
		inf:  nsInf,
		stop: infCancel,
		wait: waitWg.Wait,
	}
}

func (g *WatchGraph) handleNamespaceStatus(ctx context.Context, ns nsWithStatus) {
	switch g.ns[ns.name] { // switch on current status
	case nsDeleted:
		switch ns.status {
		case nsDeleted: // del -> del = do nothing
		case nsSelected: // del -> selected
			g.ns[ns.name] = nsSelected
			g.startWatch(ctx, ns.name)
		case nsFilteredOut: // del -> filtered out
			g.ns[ns.name] = nsFilteredOut
		default:
			panic("unreachable")
		}
	case nsSelected:
		switch ns.status {
		case nsDeleted: // selected -> del
			delete(g.ns, ns.name)
			g.stopWatch(ns.name)
		case nsSelected: // selected -> selected = do nothing
		case nsFilteredOut: // selected -> filtered out
			g.ns[ns.name] = nsFilteredOut
			g.stopWatch(ns.name)
		default:
			panic("unreachable")
		}
	case nsFilteredOut:
		switch ns.status {
		case nsDeleted: // filtered out -> del
			delete(g.ns, ns.name)
		case nsSelected: // filtered out -> selected
			g.ns[ns.name] = nsSelected
			g.startWatch(ctx, ns.name)
		case nsFilteredOut: // filtered out -> filtered out = do nothing
		default:
			panic("unreachable")
		}
	default:
		panic("unreachable")
	}
}

func (g *WatchGraph) startWatch(ctx context.Context, ns string) {
	g.opts.Log.Debug("Watch namespace", logz.NamespacedName(ns))

	for gvr, info := range g.gvrInfos.infos {
		if !info.namespaced {
			continue
		}
		g.selectedGVRs[gvr][ns] = g.startInformerForNamespacedGVR(ctx, gvr, info, ns)
	}
}

func (g *WatchGraph) stopWatch(ns string) {
	g.opts.Log.Debug("Unwatch namespace", logz.NamespacedName(ns))

	waits := make([]func(), 0, len(g.selectedGVRs))
	for gvr, ns2data := range g.selectedGVRs {
		data, ok := ns2data[ns]
		if !ok { // This must be a cluster-scoped resource
			continue
		}
		g.opts.Log.Debug("Stop watching resource", logz.GVR(gvr.Value()), logz.NamespacedName(ns))
		delete(ns2data, ns)
		data.stop()
		waits = append(waits, data.wait)
	}
	for _, w := range waits {
		w()
	}
}

func (g *WatchGraph) processWorkItem(ctx context.Context, key VertexID) *Error {
	defer g.workqueue.Done(key)

	success, err := g.processObject(ctx, key)
	if err != nil {
		g.workqueue.Forget(key)
		return err
	}
	if success {
		g.workqueue.Forget(key)
	} else {
		// put back into work queue to handle transient errors
		g.workqueue.AddRateLimited(key)
	}
	return nil
}

func (g *WatchGraph) processObject(ctx context.Context, vid VertexID) (bool /*success*/, *Error) {
	obj, err := g.id2object(vid)
	if err != nil {
		return false, err
	}
	if obj == nil {
		g.deleteVertex(ctx, vid)
		return true, nil
	}
	info, ok := g.gvrInfos.infos[vid.GVR]
	if !ok {
		g.deleteVertex(ctx, vid)
		// This shouldn't happen normally. Maybe we have outdated discovery info?
		g.opts.OnWarning(NewObjectProcessingWarning(vid.GVR.Value(), vid.Namespace, vid.Name, "Unknown GVR for an existing object"))
		// Reschedule since it's an unknown resource that might be known at some point.
		return false, nil
	}
	add := true
	if info.object.ObjectSelectorExpression != nil {
		val, _, err := info.object.ObjectSelectorExpression.ContextEval(ctx, g.opts.ObjectToEvalVars(obj, vid.GVR.Value()))
		if err != nil {
			return false, &Error{
				Message: fmt.Sprintf("query.include: object selector expression failed for %s %s/%s: %v", vid.GVR.Value(), vid.Namespace, vid.Name, err),
				Code:    InvalidArgument,
			}
		}
		add = bool(val.(types.Bool))
	}
	if !add { // Delete
		g.deleteVertex(ctx, vid)
		return true, nil
	}
	// Add
	vx := wgObject2vertex{
		gvr:  vid.GVR.Value(),
		info: info,
	}
	vd, arcsGK, gWarns, gErr := vx.Get(obj)
	for _, gWarn := range gWarns {
		g.opts.OnWarning(gWarn)
	}
	if gErr != nil {
		return false, gErr
	}
	success, arcs := g.convertArcs(vid, arcsGK)
	res := g.opts.Graph.SetVertex(ctx, vid, vd, arcs)
	if res == VertexAdded {
		// Vertex was added. Let's see if there are any pre-existing inbound arcs to it.
		// Reprocess the vertices that have arcs pointing at the newly added vertex.
		// Those arcs have "vertex does not exist" attributes set, which we need to clear.
		g.reprocessInboundFor(vid)
	}
	return success, nil
}

func (g *WatchGraph) deleteVertex(ctx context.Context, vid VertexID) {
	deleted := g.opts.Graph.DeleteVertex(ctx, vid)
	if !deleted {
		return
	}
	// Reprocess the vertices that have arcs pointing at the deleted vertex.
	// Those arcs should have "vertex does not exist" attributes set.
	g.reprocessInboundFor(vid)
}

func (g *WatchGraph) reprocessInboundFor(vid VertexID) {
	for arc := range g.opts.Graph.InboundArcsFor(vid) {
		// There may be more than one inbound arc from a to b because they can be of different types.
		// This means a single vertex may be added to the workqueue more than once.
		// This is ok because:
		// - It is quite unlikely.
		// - Workqueue will deduplicate equal work items.
		// - Duplicate processing will do no harm.
		// So, rather than doing deduplication here, we just rely on workqueue for simplicity.
		g.workqueue.Add(VertexID{
			GVR:       arc.To.GVR,
			Namespace: arc.To.Namespace,
			Name:      arc.To.Name,
		})
	}
}

func (g *WatchGraph) convertArcs(vid VertexID, arcsGK ArcSetWithData[vertexGKID, ArcType, ArcAttrs]) (bool, ArcSetWithData[VertexID, ArcType, ArcAttrs]) {
	if len(arcsGK) == 0 { // really want to reduce allocations here - do not allocate empty maps.
		return true, nil
	}
	success := true
	arcs := make(ArcSetWithData[VertexID, ArcType, ArcAttrs], len(arcsGK))
	for arc2id, arcAttrs := range arcsGK {
		info, ok := g.gvrInfos.gk2info[arc2id.To.GK]
		if !ok {
			if g.gvrInfos.gks.Has(arc2id.To.GK) {
				continue // Known GK, but it was not selected. Ignore arc.
			}
			success = false
			msg := fmt.Sprintf("Unknown GK in a reference (gk=%s, name=%s)", arc2id.To.GK, arc2id.To.Name)
			g.opts.OnWarning(NewObjectProcessingWarning(vid.GVR.Value(), vid.Namespace, vid.Name, msg))
			continue // unknown GK, reschedule
		}
		ns := metav1.NamespaceNone
		if info.namespaced { // if it's a ref to a namespaced object
			if arc2id.To.NamespaceOverride != "" {
				ns = arc2id.To.NamespaceOverride
			} else {
				if vid.Namespace == "" { // this shouldn't happen, this is an assertion to catch bugs.
					// Regardless of the reason, we don't want to use an empty namespace for a namespaced object, so... let's have this assertion.
					msg := fmt.Sprintf("Cannot determine namespace for an object (gk=%s, name=%s) referenced from a cluster-scoped object", arc2id.To.GK, arc2id.To.Name)
					g.opts.OnWarning(NewObjectProcessingWarning(vid.GVR.Value(), vid.Namespace, vid.Name, msg))
					// We don't set success to false here because retrying wouldn't change anything. Just skip the arc.
					continue
				}
				ns = vid.Namespace // then use the namespace of the object (presume same namespace).
			}
			filteredOut := g.isFilteredOutNamespace(ns)
			if filteredOut {
				// The caller is not interested in this namespace, skip the arc.
				continue
			}
		}
		destVID := VertexID{
			GVR:       info.gvr,
			Namespace: ns,
			Name:      arc2id.To.Name,
		}
		_, isSet := g.opts.Graph.VertexData(destVID) // check the graph, not informers!
		if !isSet {
			arcAttrs.DestinationDoesNotExist = true
		}
		newArc2id := ArcToID[VertexID, ArcType]{
			To:      destVID,
			ArcType: arc2id.ArcType,
		}
		arcs[newArc2id] = arcAttrs
	}
	return success, arcs
}

func newInformer(client dynamic.Interface, gvr schema.GroupVersionResource, namespace, labelSelector, fieldSelector string) cache.SharedIndexInformer {
	nsClient := client.Resource(gvr).Namespace(namespace)
	return cache.NewSharedIndexInformerWithOptions(
		&cache.ListWatch{
			ListWithContextFunc: func(ctx context.Context, opts metav1.ListOptions) (runtime.Object, error) {
				opts.LabelSelector = labelSelector
				opts.FieldSelector = fieldSelector
				return nsClient.List(context.WithoutCancel(ctx), opts) // Watch will be closed via Stop()
			},
			WatchFuncWithContext: func(ctx context.Context, opts metav1.ListOptions) (watch.Interface, error) {
				opts.LabelSelector = labelSelector
				opts.FieldSelector = fieldSelector
				return nsClient.Watch(context.WithoutCancel(ctx), opts) // Watch will be closed via Stop()
			},
		},
		&unstructured.Unstructured{},
		cache.SharedIndexInformerOptions{
			ObjectDescription: gvr.String(),
		},
	)
}
