package watch_graph //nolint:staticcheck

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

	"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/schema"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/cli-runtime/pkg/resource"
	"k8s.io/client-go/discovery"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/workqueue"
)

type Opts struct {
	Log              *slog.Logger
	Queries          []any // Elements are QueryInclude or QueryExclude
	Namespaces       Namespaces
	DiscoClient      discovery.AggregatedDiscoveryInterface
	Client           dynamic.Interface
	NewBuilder       func() *resource.Builder
	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 { //nolint: gocritic
	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.startInformerForGVR(ctx, gvr, info, metav1.NamespaceAll)
			} else {
				for ns, status := range g.ns {
					if status != nsSelected {
						continue
					}
					ns2inf[ns] = g.startInformerForGVR(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) startInformerForGVR(ctx context.Context, gvr unique.Handle[schema.GroupVersionResource], info gvrInfo, ns string) selectedGVR {
	g.opts.Log.Debug("Starting informer", logz.GVR(gvr.Value()), logz.K8sObjectNamespace(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.K8sObjectNamespace(ns))

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

func (g *WatchGraph) stopWatch(ns string) {
	g.opts.Log.Debug("Unwatch namespace", logz.K8sObjectNamespace(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.K8sObjectNamespace(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 := newObject2vertex(vid.GVR.Value(), info, g.opts.NewBuilder)
	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
}
