package it

import (
	"context"
	"encoding/json"
	"fmt"
	"testing"
	"time"

	"github.com/coder/websocket"
	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kubernetes_api/rpc"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testhelpers"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/pkg/agentcfg"
	"google.golang.org/protobuf/encoding/protojson"
	appsv1 "k8s.io/api/apps/v1"
	autoscalingv1 "k8s.io/api/autoscaling/v1"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	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/utils/ptr"
)

var (
	podGVR             = corev1.SchemeGroupVersion.WithResource("pods")
	configMapGVR       = corev1.SchemeGroupVersion.WithResource("configmaps")
	serviceaccountsGVR = corev1.SchemeGroupVersion.WithResource("serviceaccounts")

	deploymentGVR = appsv1.SchemeGroupVersion.WithResource("deployments")
	replicaSetGVR = appsv1.SchemeGroupVersion.WithResource("replicasets")

	roleGVR        = rbacv1.SchemeGroupVersion.WithResource("roles")
	roleBindingGVR = rbacv1.SchemeGroupVersion.WithResource("rolebindings")
	clusterRoleGVR = rbacv1.SchemeGroupVersion.WithResource("clusterroles")
)

func TestGraphAPIManual(t *testing.T) {
	t.SkipNow()

	const (
		kasKubernetesAPIURL = "http://gdk.test:8154/-/k8s-proxy"
		credential          = "Bearer pat:3:tL1uN9hQTcZgJ2yUpTKM" // Oh no a real GDK PAT!
	)

	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()
	conn, err := dialWS(ctx, kasKubernetesAPIURL+"/graph", credential, "gitlab-agent-graph-api")
	require.NoError(t, err)

	gw := graphWatcher{
		t:       t,
		conn:    conn,
		require: require.New(t),
		assert:  assert.New(t),
	}
	defer gw.Close()

	reqData, err := protojson.Marshal(&rpc.WatchGraphWebSocketRequest{
		Queries: []*rpc.Query{
			{
				Query: &rpc.Query_Include_{
					Include: &rpc.Query_Include{
						ResourceSelectorExpression: "group == '' && version == 'v1' && resource == 'pods'",
						Object: &rpc.Query_IncludeObject{
							ObjectSelectorExpression: "'k8s-app' in labels && labels['k8s-app'] == 'kube-dns'",
							JsonPath:                 "['metadata.name', 'metadata.namespace']",
						},
					},
				},
			},
		},
		Namespaces: &rpc.Namespaces{
			Names: []string{"kube-system"},
		},
	})
	require.NoError(t, err)

	err = conn.Write(ctx, websocket.MessageText, reqData)
	require.NoError(t, err)

	for {
		//ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		a := gw.NextAction(ctx)
		aBytes, err := json.Marshal(a)
		require.NoError(t, err)
		t.Log(string(aBytes))
		//cancel()
	}
}

func (s *integrationSuite) TestGraphAPI() {
	gitalySrv := s.startGitaly()
	s.setupGitLabMocks(gitalySrv, true)
	s.setupGitalyMocksForAgentConfig(&agentcfg.ConfigurationFile{
		Observability: &agentcfg.ObservabilityCF{
			Logging: &agentcfg.LoggingCF{
				Level:     agentcfg.LogLevelEnum_debug,
				GrpcLevel: ptr.To(agentcfg.LogLevelEnum_debug),
			},
		},
	})

	gitLabSrv := s.startGitLab()

	kas := s.startKAS(kasOptions{
		name:       "kas",
		cfg:        s.baseKASConfig(gitLabSrv),
		gitalyPort: gitalySrv.port,
	})

	s.startAgentk(agentkOptions{
		kasURL: kas.agentServerURL,
		token:  agentkToken1,
		name:   "agent1",
		ns:     s.ns,
	})

	cmaps := s.k8s.CoreV1().ConfigMaps(s.ns)
	pods := s.k8s.CoreV1().Pods(s.ns)
	serviceaccounts := s.k8s.CoreV1().ServiceAccounts(s.ns)
	rbacRoles := s.k8s.RbacV1().Roles(s.ns)
	rbacRoleBindings := s.k8s.RbacV1().RoleBindings(s.ns)
	deps := s.k8s.AppsV1().Deployments(s.ns)
	s.Run("create,delete,owner references,whole object,namespace by label", func() {
		gw := s.newGraphWatcher(kas, &rpc.WatchGraphWebSocketRequest{
			Queries: []*rpc.Query{
				{ // Want configmaps with a certain label, nothing else
					Query: &rpc.Query_Include_{
						Include: &rpc.Query_Include{
							ResourceSelectorExpression: "group == '' && version == 'v1' && resource == 'configmaps'",
							Object: &rpc.Query_IncludeObject{
								LabelSelector: "ls=test1",
							},
						},
					},
				},
			},
			Namespaces: &rpc.Namespaces{
				LabelSelector: "ls=" + nsLabel, // want them in any namespace with this label
			},
		})
		defer gw.Close()

		// 1. Setup ConfigMap 1
		cm1 := &corev1.ConfigMap{
			TypeMeta: metav1.TypeMeta{
				Kind:       "ConfigMap",
				APIVersion: "v1",
			},
			ObjectMeta: metav1.ObjectMeta{
				Name:      "cm1",
				Namespace: s.ns,
				Labels: map[string]string{
					"ls": "test1",
				},
				OwnerReferences: []metav1.OwnerReference{ // this owner reference is ignored since we are not watching namespaces.
					{
						APIVersion: "v1",
						Kind:       "Namespace",
						Name:       s.ns,
						UID:        s.nsObj.UID,
						Controller: ptr.To(true),
					},
				},
			},
		}
		cm1Obj, err := cmaps.Create(s.ctx, cm1, metav1.CreateOptions{})
		s.Require().NoError(err)

		cm1Unstr := s.toUnstructured(cm1, cm1Obj)
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			g.AssertVertex(configMapGVR, cm1Obj.Namespace, cm1Obj.Name, cm1Unstr, nil, nil)
			g.AssertNoOtherVertices()
		}))

		// 2. Setup ConfigMap 2
		cm2 := &corev1.ConfigMap{
			TypeMeta: metav1.TypeMeta{
				Kind:       "ConfigMap",
				APIVersion: "v1",
			},
			ObjectMeta: metav1.ObjectMeta{
				Name:      "cm2",
				Namespace: s.ns,
				Labels: map[string]string{
					"ls": "test1",
				},
				OwnerReferences: []metav1.OwnerReference{ // valid owner reference
					{
						APIVersion: "v1",
						Kind:       "ConfigMap",
						Name:       cm1Obj.Name,
						UID:        cm1Obj.UID,
					},
				},
			},
		}
		cm2Obj, err := cmaps.Create(s.ctx, cm2, metav1.CreateOptions{})
		s.Require().NoError(err)
		cm2Unstr := s.toUnstructured(cm2, cm2Obj)
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			cm1Vx := g.AssertVertex(configMapGVR, cm1Obj.Namespace, cm1Obj.Name, cm1Unstr, nil, nil)
			g.AssertVertex(configMapGVR, cm2Obj.Namespace, cm2Obj.Name, cm2Unstr, nil, map[arcToVertex]arcAttrs{
				{To: cm1Vx, ArcType: ownerReferenceArcType}: nil,
			})
			g.AssertNoOtherVertices()
		}))

		// 3. Delete ConfigMap 2
		err = cmaps.Delete(s.ctx, cm2Obj.Name, metav1.DeleteOptions{})
		s.Require().NoError(err)
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			g.AssertVertex(configMapGVR, cm1Obj.Namespace, cm1Obj.Name, cm1Unstr, nil, nil)
			g.AssertNoOtherVertices()
		}))

		// 4. Delete ConfigMap 1
		err = cmaps.Delete(s.ctx, cm1Obj.Name, metav1.DeleteOptions{})
		s.Require().NoError(err)
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			g.AssertNoOtherVertices()
		}))
	})
	s.Run("owner references", func() {
		gw := s.newGraphWatcher(kas, &rpc.WatchGraphWebSocketRequest{
			Queries: []*rpc.Query{
				{ // Want configmaps with a certain label, nothing else
					Query: &rpc.Query_Include_{
						Include: &rpc.Query_Include{
							ResourceSelectorExpression: "group == 'rbac.authorization.k8s.io' && version == 'v1' && resource in ['roles', 'rolebindings']",
						},
					},
				},
			},
			Namespaces: &rpc.Namespaces{
				LabelSelector: "ls=" + nsLabel, // want them in any namespace with this label
			},
		})
		defer gw.Close()

		r1 := &rbacv1.Role{
			TypeMeta: metav1.TypeMeta{
				Kind:       "Role",
				APIVersion: "rbac.authorization.k8s.io/v1",
			},
			ObjectMeta: metav1.ObjectMeta{
				Name:      "role1",
				Namespace: s.ns,
			},
		}
		rb1 := &rbacv1.RoleBinding{
			TypeMeta: metav1.TypeMeta{
				Kind:       "RoleBinding",
				APIVersion: "rbac.authorization.k8s.io/v1",
			},
			ObjectMeta: metav1.ObjectMeta{
				Name:      "rb1",
				Namespace: s.ns,
			},
			RoleRef: rbacv1.RoleRef{
				APIGroup: "rbac.authorization.k8s.io",
				Kind:     "Role",
				Name:     r1.Name,
			},
		}
		// 1. Create a RoleBinding that refs a Role that does not exist.
		rb1obj, err := rbacRoleBindings.Create(s.ctx, rb1, metav1.CreateOptions{})
		s.Require().NoError(err)
		rb1Unstr := s.toUnstructured(rb1, rb1obj)
		r1Vx := newVertex(roleGVR, r1.Namespace, r1.Name)
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			g.AssertVertex(roleBindingGVR, rb1obj.Namespace, rb1obj.Name, rb1Unstr, nil, map[arcToVertex]arcAttrs{
				{To: r1Vx, ArcType: referenceArcType}: {"e": true},
			})
			g.AssertNoOtherVertices()
		}))

		// 2. Create the Role that is referenced by the RoleBinding.
		r1obj, err := rbacRoles.Create(s.ctx, r1, metav1.CreateOptions{})
		s.Require().NoError(err)
		r1Unstr := s.toUnstructured(r1, r1obj)
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			r1Vx = g.AssertVertex(roleGVR, r1obj.Namespace, r1obj.Name, r1Unstr, nil, nil)
			g.AssertVertex(roleBindingGVR, rb1obj.Namespace, rb1obj.Name, rb1Unstr, nil, map[arcToVertex]arcAttrs{
				{To: r1Vx, ArcType: referenceArcType}: nil,
			})
			g.AssertNoOtherVertices()
		}))

		// 3. Delete the Role again.
		err = rbacRoles.Delete(s.ctx, r1.Name, metav1.DeleteOptions{})
		s.Require().NoError(err)

		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			g.AssertVertex(roleBindingGVR, rb1obj.Namespace, rb1obj.Name, rb1Unstr, nil, map[arcToVertex]arcAttrs{
				{To: r1Vx, ArcType: referenceArcType}: {"e": true},
			})
			g.AssertNoOtherVertices()
		}))
	})
	s.Run("create,JSON path,namespace by name", func() {
		gw := s.newGraphWatcher(kas, &rpc.WatchGraphWebSocketRequest{
			Queries: []*rpc.Query{
				{ // Want configmaps with a certain label, nothing else
					Query: &rpc.Query_Include_{
						Include: &rpc.Query_Include{
							ResourceSelectorExpression: "group == '' && version == 'v1' && resource == 'configmaps'",
							Object: &rpc.Query_IncludeObject{
								LabelSelector: "ls=test2",
								JsonPath:      "['metadata.name', 'metadata.namespace']", // only want these fields
							},
						},
					},
				},
			},
			Namespaces: &rpc.Namespaces{
				Names: []string{s.ns}, // only this namespace
			},
		})
		defer gw.Close()

		cm1 := &corev1.ConfigMap{
			TypeMeta: metav1.TypeMeta{
				Kind:       "ConfigMap",
				APIVersion: "v1",
			},
			ObjectMeta: metav1.ObjectMeta{
				Name:      "cm1",
				Namespace: s.ns,
				Labels: map[string]string{
					"ls": "test2",
				},
			},
		}
		cm1Obj, err := cmaps.Create(s.ctx, cm1, metav1.CreateOptions{})
		s.Require().NoError(err)
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			g.AssertVertex(configMapGVR, cm1Obj.Namespace, cm1Obj.Name, nil, []any{cm1Obj.Name, cm1Obj.Namespace}, nil)
			g.AssertNoOtherVertices()
		}))
	})
	s.Run("roots - vertex create, delete", func() {
		gw := s.newGraphWatcher(kas, &rpc.WatchGraphWebSocketRequest{
			Queries: []*rpc.Query{
				{
					Query: &rpc.Query_Include_{
						Include: &rpc.Query_Include{
							ResourceSelectorExpression: "resource in ['deployments', 'replicasets', 'pods']",
						},
					},
				},
			},
			Namespaces: &rpc.Namespaces{
				LabelSelector: "ls=" + nsLabel, // want them in any namespace with this label
			},
			Roots: &rpc.Roots{
				ObjectSelectorExpressions: []string{"group == '' && version == 'v1' && resource == 'pods' && 'is-root1' in labels"},
			},
		})
		defer gw.Close()

		s.T().Log("1. Create a Deployment")
		d1 := &appsv1.Deployment{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "d1",
				Namespace: s.ns,
			},
			Spec: appsv1.DeploymentSpec{
				Replicas: ptr.To[int32](2), // 2 to test various edge cases, 1 is not enough.
				Selector: &metav1.LabelSelector{
					MatchLabels: map[string]string{
						"ls": "test-roots",
					},
				},
				Template: corev1.PodTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							"ls":       "test-roots",
							"is-root1": "true",
						},
					},
					Spec: corev1.PodSpec{
						Containers: []corev1.Container{
							{
								Name:            "c1",
								Image:           "alpine",
								Args:            []string{"sleep", "10m"},
								ImagePullPolicy: corev1.PullIfNotPresent,
							},
						},
						TerminationGracePeriodSeconds: ptr.To[int64](1),
					},
				},
			},
		}

		_, err := deps.Create(s.ctx, d1, metav1.CreateOptions{})
		s.Require().NoError(err)

		s.T().Log("2. Wait for the graph to reach the expected state")
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			depl := g.AssertOne(deploymentGVR, s.ns, nil)
			rs := g.AssertOne(replicaSetGVR, s.ns, map[arcToVertex]arcAttrs{
				{To: depl, ArcType: ownerReferenceArcType}: {"c": true, "b": true},
			})
			g.AssertEach(podGVR, s.ns, 2, map[arcToVertex]arcAttrs{
				{To: rs, ArcType: ownerReferenceArcType}: {"c": true, "b": true},
			})
			g.AssertNoOtherVertices()
		}))

		s.T().Log("3. Scale the deployment to zero. There will be no root nodes (no Pods) but ReplicaSet and Deployment will still be there. They are unreachable now (no roots), and that's what we expect")
		_, err = deps.UpdateScale(s.ctx, d1.Name, &autoscalingv1.Scale{
			ObjectMeta: metav1.ObjectMeta{
				Name: d1.Name,
			},
			Spec: autoscalingv1.ScaleSpec{
				Replicas: 0,
			},
		}, metav1.UpdateOptions{})
		s.Require().NoError(err)

		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			g.AssertNoOtherVertices()
		}))

		s.T().Log("4. Scale the deployment back to 1")
		_, err = deps.UpdateScale(s.ctx, d1.Name, &autoscalingv1.Scale{
			ObjectMeta: metav1.ObjectMeta{
				Name: d1.Name,
			},
			Spec: autoscalingv1.ScaleSpec{
				Replicas: 1,
			},
		}, metav1.UpdateOptions{})
		s.Require().NoError(err)

		s.T().Log("5. Wait for the graph to reach the expected state - expect to see a Pod, a ReplicaSet, and a Deployment")
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			depl := g.AssertOne(deploymentGVR, s.ns, nil)
			rs := g.AssertOne(replicaSetGVR, s.ns, map[arcToVertex]arcAttrs{
				{To: depl, ArcType: ownerReferenceArcType}: {"c": true, "b": true},
			})
			g.AssertEach(podGVR, s.ns, 1, map[arcToVertex]arcAttrs{
				{To: rs, ArcType: ownerReferenceArcType}: {"c": true, "b": true},
			})
			g.AssertNoOtherVertices()
		}))

		s.T().Log("6. Delete the Deployment")
		err = deps.Delete(s.ctx, d1.Name, metav1.DeleteOptions{})
		s.Require().NoError(err)

		s.T().Log("7. Wait for the graph to reach the expected state - expect to see everything gone. All that shines turns to rust. All the stands in time turns to dust")
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			g.AssertNoOtherVertices()
		}))
	})
	s.Run("roots - arc create, delete", func() {
		gw := s.newGraphWatcher(kas, &rpc.WatchGraphWebSocketRequest{
			Queries: []*rpc.Query{
				{
					Query: &rpc.Query_Include_{
						Include: &rpc.Query_Include{
							ResourceSelectorExpression: "resource == 'configmaps'",
						},
					},
				},
			},
			Namespaces: &rpc.Namespaces{
				LabelSelector: "ls=" + nsLabel, // want them in any namespace with this label
			},
			Roots: &rpc.Roots{
				ObjectSelectorExpressions: []string{"group == '' && resource == 'configmaps' && 'is-root2' in labels"},
			},
		})
		defer gw.Close()

		s.T().Log("1. Setup ConfigMap 1")
		cm1 := &corev1.ConfigMap{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "rootcm1",
				Namespace: s.ns,
			},
		}
		cm1Obj, err := cmaps.Create(s.ctx, cm1, metav1.CreateOptions{})
		s.Require().NoError(err)

		s.T().Log("2. Setup ConfigMap 2")
		cm2 := &corev1.ConfigMap{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "rootcm2",
				Namespace: s.ns,
				Labels: map[string]string{
					"is-root2": "true",
				},
				OwnerReferences: []metav1.OwnerReference{
					{
						APIVersion: "v1",
						Kind:       "ConfigMap",
						Name:       cm1Obj.Name,
						UID:        cm1Obj.UID,
					},
				},
			},
		}
		_, err = cmaps.Create(s.ctx, cm2, metav1.CreateOptions{})
		s.Require().NoError(err)

		s.T().Log("3. Wait for the graph to reach the expected state - two ConfigMaps linked with an arc")
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			cm1Vid := g.AssertVertex(configMapGVR, s.ns, cm1.Name, nil, nil, nil)
			g.AssertVertex(configMapGVR, s.ns, cm2.Name, nil, nil, map[arcToVertex]arcAttrs{
				{To: cm1Vid, ArcType: ownerReferenceArcType}: nil,
			})
			g.AssertNoOtherVertices()
		}))

		s.T().Log("4. Remove ConfigMap 2 to ConfigMaps 1 arc - remove the owner reference")
		cm2.OwnerReferences = nil
		_, err = cmaps.Update(s.ctx, cm2, metav1.UpdateOptions{})
		s.Require().NoError(err)

		s.T().Log("5. Wait for the graph to reach the expected state - a single ConfigMap")
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			g.AssertVertex(configMapGVR, s.ns, cm2.Name, nil, nil, nil)
			g.AssertNoOtherVertices()
		}))
	})
	s.Run("roots - delete reachable vertex", func() {
		gw := s.newGraphWatcher(kas, &rpc.WatchGraphWebSocketRequest{
			Queries: []*rpc.Query{
				{
					Query: &rpc.Query_Include_{
						Include: &rpc.Query_Include{
							ResourceSelectorExpression: "resource == 'configmaps' || resource == 'pods'",
						},
					},
				},
			},
			Namespaces: &rpc.Namespaces{
				LabelSelector: "ls=" + nsLabel, // want them in any namespace with this label
			},
			Roots: &rpc.Roots{
				ObjectSelectorExpressions: []string{"group == '' && resource == 'pods' && 'is-root3' in labels"},
			},
		})
		defer gw.Close()

		s.T().Log("1. Setup ConfigMap")
		cm4 := &corev1.ConfigMap{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "cm4",
				Namespace: s.ns,
			},
		}
		_, err := cmaps.Create(s.ctx, cm4, metav1.CreateOptions{})
		s.Require().NoError(err)

		s.T().Log("2. Setup Pods")
		p1 := &corev1.Pod{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "rootpod1",
				Namespace: s.ns,
				Labels: map[string]string{
					"is-root3": "true",
				},
			},
			Spec: corev1.PodSpec{
				Volumes: []corev1.Volume{
					{
						Name: "configmap",
						VolumeSource: corev1.VolumeSource{
							ConfigMap: &corev1.ConfigMapVolumeSource{
								LocalObjectReference: corev1.LocalObjectReference{
									Name: cm4.Name,
								},
							},
						},
					},
				},
				Containers: []corev1.Container{
					{
						Name:            "c1",
						Image:           "alpine",
						Args:            []string{"sleep", "10m"},
						ImagePullPolicy: corev1.PullIfNotPresent,
					},
				},
				TerminationGracePeriodSeconds: ptr.To[int64](1),
			},
		}
		p2 := p1.DeepCopy()
		p2.Name = "rootpod2"

		_, err = pods.Create(s.ctx, p1, metav1.CreateOptions{})
		s.Require().NoError(err)
		_, err = pods.Create(s.ctx, p2, metav1.CreateOptions{})
		s.Require().NoError(err)

		s.T().Log("3. Wait for the graph to reach the expected state - Pods and ConfigMap linked with arcs")
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			cm4Vid := g.AssertVertex(configMapGVR, s.ns, cm4.Name, nil, nil, nil)
			rootCAVid := g.AssertVertex(configMapGVR, s.ns, "kube-root-ca.crt", nil, nil, nil)
			g.AssertVertex(podGVR, s.ns, p1.Name, nil, nil, map[arcToVertex]arcAttrs{
				{To: cm4Vid, ArcType: referenceArcType}:    nil,
				{To: rootCAVid, ArcType: referenceArcType}: nil,
			})
			g.AssertVertex(podGVR, s.ns, p2.Name, nil, nil, map[arcToVertex]arcAttrs{
				{To: cm4Vid, ArcType: referenceArcType}:    nil,
				{To: rootCAVid, ArcType: referenceArcType}: nil,
			})
			g.AssertNoOtherVertices()
		}))

		s.T().Log("4. Remove the ConfigMap")
		err = cmaps.Delete(s.ctx, cm4.Name, metav1.DeleteOptions{})
		s.Require().NoError(err)

		s.T().Log("5. Wait for the graph to reach the expected state - a Pod")
		cm4Vid := newVertex(configMapGVR, s.ns, cm4.Name)
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			rootCAVid := g.AssertVertex(configMapGVR, s.ns, "kube-root-ca.crt", nil, nil, nil)
			g.AssertVertex(podGVR, s.ns, p1.Name, nil, nil, map[arcToVertex]arcAttrs{
				{To: cm4Vid, ArcType: referenceArcType}:    {"e": true},
				{To: rootCAVid, ArcType: referenceArcType}: nil,
			})
			g.AssertVertex(podGVR, s.ns, p2.Name, nil, nil, map[arcToVertex]arcAttrs{
				{To: cm4Vid, ArcType: referenceArcType}:    {"e": true},
				{To: rootCAVid, ArcType: referenceArcType}: nil,
			})
			g.AssertNoOtherVertices()
		}))
	})
	s.Run("roots - undirected", func() {
		gw := s.newGraphWatcher(kas, &rpc.WatchGraphWebSocketRequest{
			Queries: []*rpc.Query{
				{
					Query: &rpc.Query_Include_{
						Include: &rpc.Query_Include{
							ResourceSelectorExpression: "resource in ['deployments', 'replicasets', 'pods', 'serviceaccounts', 'rolebindings', 'clusterroles']",
						},
					},
				},
			},
			Namespaces: &rpc.Namespaces{
				LabelSelector: "ls=" + nsLabel, // want them in any namespace with this label
			},
			Roots: &rpc.Roots{
				ObjectSelectorExpressions: []string{"group == 'apps' && resource == 'deployments' && 'is-root4' in labels"},
				IgnoreArcDirection:        []string{string(ownerReferenceArcType), string(referenceArcType), string(transitiveReferenceArcType)},
			},
		})
		defer gw.Close()

		s.T().Log("1. Create a ServiceAccount")
		sa1 := &corev1.ServiceAccount{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "root4-sa",
				Namespace: s.ns,
			},
		}
		_, err := serviceaccounts.Create(s.ctx, sa1, metav1.CreateOptions{})
		s.Require().NoError(err)

		s.T().Log("2. Create a RoleBinding")
		rb2 := &rbacv1.RoleBinding{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "rb2",
				Namespace: s.ns,
			},
			Subjects: []rbacv1.Subject{
				{
					Kind: "ServiceAccount",
					Name: sa1.Name,
				},
			},
			RoleRef: rbacv1.RoleRef{
				APIGroup: "rbac.authorization.k8s.io",
				Kind:     "ClusterRole",
				Name:     "view",
			},
		}
		_, err = rbacRoleBindings.Create(s.ctx, rb2, metav1.CreateOptions{})
		s.Require().NoError(err)

		s.T().Log("3. Create a Deployment")
		d1 := &appsv1.Deployment{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "d4",
				Namespace: s.ns,
				Labels: map[string]string{
					"is-root4": "true",
				},
			},
			Spec: appsv1.DeploymentSpec{
				Replicas: ptr.To[int32](2), // 2 to test various edge cases, 1 is not enough.
				Selector: &metav1.LabelSelector{
					MatchLabels: map[string]string{
						"ls": "test-roots4",
					},
				},
				Template: corev1.PodTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							"ls": "test-roots4",
						},
					},
					Spec: corev1.PodSpec{
						Containers: []corev1.Container{
							{
								Name:            "c1",
								Image:           "alpine",
								Args:            []string{"sleep", "10m"},
								ImagePullPolicy: corev1.PullIfNotPresent,
							},
						},
						TerminationGracePeriodSeconds: ptr.To[int64](1),
						ServiceAccountName:            sa1.Name,
					},
				},
			},
		}

		_, err = deps.Create(s.ctx, d1, metav1.CreateOptions{})
		s.Require().NoError(err)

		s.T().Log("4. Wait for the graph to reach the expected state")
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			crView := g.AssertVertex(clusterRoleGVR, "", "view", nil, nil, nil)
			sa := g.AssertVertex(serviceaccountsGVR, s.ns, sa1.Name, nil, nil, nil)
			g.AssertVertex(roleBindingGVR, s.ns, rb2.Name, nil, nil, map[arcToVertex]arcAttrs{
				{To: crView, ArcType: referenceArcType}: nil,
				{To: sa, ArcType: referenceArcType}:     nil,
			})
			depl := g.AssertOne(deploymentGVR, s.ns, map[arcToVertex]arcAttrs{
				{To: sa, ArcType: transitiveReferenceArcType}: nil,
			})
			rs := g.AssertOne(replicaSetGVR, s.ns, map[arcToVertex]arcAttrs{
				{To: depl, ArcType: ownerReferenceArcType}:    {"c": true, "b": true},
				{To: sa, ArcType: transitiveReferenceArcType}: nil,
			})
			g.AssertEach(podGVR, s.ns, 2, map[arcToVertex]arcAttrs{
				{To: rs, ArcType: ownerReferenceArcType}: {"c": true, "b": true},
				{To: sa, ArcType: referenceArcType}:      nil,
			})
			g.AssertNoOtherVertices()
		}))

		s.T().Log("5. Scale the deployment to zero. There will be no root nodes (no Pods) but ReplicaSet and Deployment will still be there.")
		_, err = deps.UpdateScale(s.ctx, d1.Name, &autoscalingv1.Scale{
			ObjectMeta: metav1.ObjectMeta{
				Name: d1.Name,
			},
			Spec: autoscalingv1.ScaleSpec{
				Replicas: 0,
			},
		}, metav1.UpdateOptions{})
		s.Require().NoError(err)

		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			crView := g.AssertVertex(clusterRoleGVR, "", "view", nil, nil, nil)
			sa := g.AssertVertex(serviceaccountsGVR, s.ns, sa1.Name, nil, nil, nil)
			g.AssertVertex(roleBindingGVR, s.ns, rb2.Name, nil, nil, map[arcToVertex]arcAttrs{
				{To: crView, ArcType: referenceArcType}: nil,
				{To: sa, ArcType: referenceArcType}:     nil,
			})
			depl := g.AssertOne(deploymentGVR, s.ns, map[arcToVertex]arcAttrs{
				{To: sa, ArcType: transitiveReferenceArcType}: nil,
			})
			g.AssertOne(replicaSetGVR, s.ns, map[arcToVertex]arcAttrs{
				{To: depl, ArcType: ownerReferenceArcType}:    {"c": true, "b": true},
				{To: sa, ArcType: transitiveReferenceArcType}: nil,
			})
			g.AssertNoOtherVertices()
		}))

		s.T().Log("6. Scale the deployment back to 1")
		_, err = deps.UpdateScale(s.ctx, d1.Name, &autoscalingv1.Scale{
			ObjectMeta: metav1.ObjectMeta{
				Name: d1.Name,
			},
			Spec: autoscalingv1.ScaleSpec{
				Replicas: 1,
			},
		}, metav1.UpdateOptions{})
		s.Require().NoError(err)

		s.T().Log("7. Wait for the graph to reach the expected state - expect to see a Pod, a ReplicaSet, and a Deployment")
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			crView := g.AssertVertex(clusterRoleGVR, "", "view", nil, nil, nil)
			sa := g.AssertVertex(serviceaccountsGVR, s.ns, sa1.Name, nil, nil, nil)
			g.AssertVertex(roleBindingGVR, s.ns, rb2.Name, nil, nil, map[arcToVertex]arcAttrs{
				{To: crView, ArcType: referenceArcType}: nil,
				{To: sa, ArcType: referenceArcType}:     nil,
			})
			depl := g.AssertOne(deploymentGVR, s.ns, map[arcToVertex]arcAttrs{
				{To: sa, ArcType: transitiveReferenceArcType}: nil,
			})
			rs := g.AssertOne(replicaSetGVR, s.ns, map[arcToVertex]arcAttrs{
				{To: depl, ArcType: ownerReferenceArcType}:    {"c": true, "b": true},
				{To: sa, ArcType: transitiveReferenceArcType}: nil,
			})
			g.AssertEach(podGVR, s.ns, 1, map[arcToVertex]arcAttrs{
				{To: rs, ArcType: ownerReferenceArcType}: {"c": true, "b": true},
				{To: sa, ArcType: referenceArcType}:      nil,
			})
			g.AssertNoOtherVertices()
		}))

		s.T().Log("8. Delete the Deployment")
		err = deps.Delete(s.ctx, d1.Name, metav1.DeleteOptions{})
		s.Require().NoError(err)

		s.T().Log("9. Wait for the graph to reach the expected state - expect to see everything gone. All that shines turns to rust. All the stands in time turns to dust")
		gw.WatchUntil(s.ctx, s.assertGraphUntil(func(g *assertGraph) {
			g.AssertNoOtherVertices()
		}))
	})
}

type assertGraphPanicVal int

const apv assertGraphPanicVal = -23

type assertGraph struct {
	t       *testing.T
	g       *objectGraph
	visited sets.Set[vertex]
}

func (g *assertGraph) AssertOne(gvr schema.GroupVersionResource, ns string, expectedArcSet map[arcToVertex]arcAttrs) vertex {
	vds := g.AssertEach(gvr, ns, 1, expectedArcSet)
	for vid := range vds {
		return vid
	}
	panic("unreachable")
}

func (g *assertGraph) AssertVertex(gvr schema.GroupVersionResource, ns, name string,
	expectedVertexObject map[string]any,
	expectedVertexJSONPath []any,
	expectedArcSet map[arcToVertex]arcAttrs) vertex {
	vid := newVertex(gvr, ns, name)
	actualData, isSet := g.g.ID2vertex[vid]
	if !isSet {
		g.t.Logf("vid=%s: MISS: no vertex", vid)
		panic(apv)
	}
	g.t.Logf("vid=%s: MATCH: got vertex", vid)
	switch {
	case expectedVertexObject != nil:
		diff := cmp.Diff(expectedVertexObject, actualData.Object)
		if diff != "" {
			g.t.Logf("vid=%s: MISS: vertex data:\n%s", vid, diff)
			panic(apv)
		}
		g.t.Logf("vid=%s: MATCH: vertex data", vid)
	case expectedVertexJSONPath != nil:
		diff := cmp.Diff(expectedVertexJSONPath, actualData.JSONPath)
		if diff != "" {
			g.t.Logf("vid=%s: MISS: vertex JSON path data:\n%s", vid, diff)
			panic(apv)
		}
		g.t.Logf("vid=%s: MATCH: vertex JSON path data", vid)
	}
	actualArcSet := g.g.ID2arcSet[vid]
	diff := cmp.Diff(expectedArcSet, actualArcSet)
	if diff != "" {
		g.t.Logf("vid=%s: MISS: arcs:\n%s", vid, diff)
		panic(apv)
	}
	g.t.Logf("vid=%s: MATCH: all arcs", vid)
	g.visited.Insert(vid)
	return vid
}

func (g *assertGraph) AssertEach(gvr schema.GroupVersionResource, ns string, expectedCount int, expectedArcSet map[arcToVertex]arcAttrs) map[vertex]vertexData {
	vds := g.DataFor(gvr, ns)
	if len(vds) != expectedCount {
		g.t.Logf("%s/%s/%s, ns=%s: MISS: expected %d, got: %d", gvr.Group, gvr.Version, gvr.Resource, ns, expectedCount, len(vds))
		panic(apv)
	}
	g.t.Logf("%s/%s/%s, ns=%s: MATCH: got %d", gvr.Group, gvr.Version, gvr.Resource, ns, expectedCount)
	for vid := range vds {
		actualArcSet := g.g.ID2arcSet[vid]
		diff := cmp.Diff(expectedArcSet, actualArcSet)
		if diff != "" {
			g.t.Logf("vid=%s: MISS: arcs:\n%s", vid, diff)
			panic(apv)
		}
		g.visited.Insert(vid)
	}
	g.t.Logf("%s/%s/%s, ns=%s: MATCH: all arcs", gvr.Group, gvr.Version, gvr.Resource, ns)
	return vds
}

func (g *assertGraph) DataFor(gvr schema.GroupVersionResource, ns string) map[vertex]vertexData {
	res := make(map[vertex]vertexData)

	for vid, vd := range g.g.ID2vertex {
		if vid.GVR() != gvr || vid.Namespace != ns {
			continue
		}
		res[vid] = vd
	}

	return res
}

func (g *assertGraph) AssertNoOtherVertices() {
	gVids := sets.Set[vertex]{}
	for vid := range g.g.ID2vertex {
		gVids.Insert(vid)
	}
	diff := cmp.Diff(gVids, g.visited)
	if diff != "" {
		g.t.Logf("Have not asserted all verices:\n%s", diff)
		panic(apv)
	}
}

func (s *integrationSuite) assertGraphUntil(f func(*assertGraph)) func(*objectGraph) bool {
	return func(g *objectGraph) (ret bool) {
		a := &assertGraph{
			t:       s.T(),
			g:       g,
			visited: sets.Set[vertex]{},
		}
		defer func() {
			r := recover()
			switch r {
			case apv:
				ret = false
			case nil:
			default:
				panic(r)
			}
		}()
		s.T().Log("Checking graph assertions")
		f(a)
		s.T().Log("All graph assertions passed")
		return true
	}
}

func (s *integrationSuite) newGraphWatcher(kas *kasHolder, req *rpc.WatchGraphWebSocketRequest) *graphWatcher {
	reqData, err := protojson.Marshal(req)
	s.Require().NoError(err)

	conn := s.dialWS(kas, testhelpers.AgentkKey1)

	err = conn.Write(s.ctx, websocket.MessageText, reqData)
	if err != nil {
		_ = conn.CloseNow()
	}
	s.Require().NoError(err)

	return &graphWatcher{
		t:       s.T(),
		conn:    conn,
		require: s.Require(),
		assert:  s.Assert(),
		graph:   newObjectGraph(),
	}
}

func (s *integrationSuite) toUnstructured(src, o runtime.Object) map[string]any {
	ou, err := runtime.DefaultUnstructuredConverter.ToUnstructured(o)
	s.Require().NoError(err)

	// Copy GVK
	u := unstructured.Unstructured{Object: ou}
	gvk := src.GetObjectKind().GroupVersionKind()
	u.SetKind(gvk.Kind)
	u.SetAPIVersion(gvk.GroupVersion().String())

	return ou
}

type graphWatcher struct {
	t       *testing.T
	conn    *websocket.Conn
	require *require.Assertions
	assert  *assert.Assertions
	graph   *objectGraph
	resp    jsonWatchGraphResponse
}

func (w *graphWatcher) Close() {
	err := w.conn.Close(websocket.StatusNormalClosure, "")
	w.assert.NoError(err)
}

func (w *graphWatcher) WatchUntil(ctx context.Context, f func(*objectGraph) bool) {
	for {
		a := w.NextAction(ctx)
		w.applyAction(a)
		if f(w.graph) {
			return
		}
	}
}

func (w *graphWatcher) NextAction(ctx context.Context) jsonWatchGraphAction {
	for {
		if len(w.resp.Actions) > 0 {
			a := w.resp.Actions[0]
			w.resp.Actions = w.resp.Actions[1:]
			return a
		}
		_, msg, err := w.conn.Read(ctx)
		w.require.NoError(err)

		err = json.Unmarshal(msg, &w.resp)
		w.require.NoError(err)
		for _, warn := range w.resp.Warnings {
			warnData, err := json.Marshal(warn)
			w.require.NoError(err)
			w.t.Log(string(warnData))
		}
	}
}

func (w *graphWatcher) applyAction(a jsonWatchGraphAction) {
	switch {
	case a.SetVertex != nil:
		sv := a.SetVertex
		w.t.Log("SetVertex", sv.Vertex)
		w.graph.SetVertex(sv.Vertex, sv.Object, sv.JSONPath)
	case a.SetArc != nil:
		sa := a.SetArc
		w.t.Log("SetArc", sa.Source, "->", sa.Destination, sa.Type)
		w.graph.SetArc(sa.Source, sa.Destination, sa.Type, sa.Attributes)
	case a.DeleteVertex != nil:
		dv := a.DeleteVertex
		w.t.Log("DeleteVertex", dv.Vertex)
		w.graph.DeleteVertex(dv.Vertex)
	case a.DeleteArc != nil:
		da := a.DeleteArc
		w.t.Log("DeleteArc", da.Source, "->", da.Destination, da.Type)
		w.graph.DeleteArc(da.Source, da.Destination, da.Type)
	default:
		w.t.Fatalf("unknown action: %T", a)
	}
}

type jsonWatchGraphResponse struct {
	Actions  []jsonWatchGraphAction  `json:"actions,omitempty"`
	Warnings []jsonWatchGraphWarning `json:"warnings,omitempty"`
}

type jsonWatchGraphAction struct {
	SetVertex    *jsonSetVertex    `json:"svx,omitempty"`
	DeleteVertex *jsonDeleteVertex `json:"dvx,omitempty"`
	SetArc       *jsonSetArc       `json:"sarc,omitempty"`
	DeleteArc    *jsonDeleteArc    `json:"darc,omitempty"`
}

type vertex struct {
	Group     string `json:"g,omitempty"`
	Version   string `json:"v"`
	Resource  string `json:"r"`
	Namespace string `json:"ns,omitempty"`
	Name      string `json:"n"`
}

func newVertex(gvr schema.GroupVersionResource, ns, name string) vertex {
	return vertex{
		Group:     gvr.Group,
		Version:   gvr.Version,
		Resource:  gvr.Resource,
		Namespace: ns,
		Name:      name,
	}
}

func (v vertex) String() string {
	return fmt.Sprintf("%s/%s/%s ns=%s n=%s", v.Group, v.Version, v.Resource, v.Namespace, v.Name)
}

func (v vertex) GVR() schema.GroupVersionResource {
	return schema.GroupVersionResource{
		Group:    v.Group,
		Version:  v.Version,
		Resource: v.Resource,
	}
}

type vertexData struct {
	Object   map[string]any
	JSONPath []any
}

type jsonSetVertex struct {
	Vertex   vertex         `json:"vx"`
	Object   map[string]any `json:"o,omitempty"`
	JSONPath []any          `json:"j,omitempty"`
}

type jsonDeleteVertex struct {
	Vertex vertex `json:"vx"`
}

type arcType string

const (
	ownerReferenceArcType      arcType = "or"
	referenceArcType           arcType = "r"
	transitiveReferenceArcType arcType = "t"
)

// arcAttrs is the arc's attributes.
// We use a map rather than a struct to make sure we catch all fields that are emitted by the server.
// That way the tests will catch any changes in the server behavior.
type arcAttrs map[string]any

type Arc struct {
	Source      vertex  `json:"s"`
	Destination vertex  `json:"d"`
	Type        arcType `json:"t"`
}

type jsonSetArc struct {
	Arc
	Attributes arcAttrs `json:"a,omitempty"`
}

type jsonDeleteArc struct {
	Arc
}

type jsonWatchGraphWarning struct {
	Type       string         `json:"t"`
	Message    string         `json:"m"`
	Attributes map[string]any `json:"a,omitempty"`
}

type arcToVertex struct {
	To      vertex
	ArcType arcType
}

type objectGraph struct {
	ID2vertex map[vertex]vertexData
	ID2arcSet map[vertex]map[arcToVertex]arcAttrs
}

func newObjectGraph() *objectGraph {
	return &objectGraph{
		ID2vertex: make(map[vertex]vertexData),
		ID2arcSet: make(map[vertex]map[arcToVertex]arcAttrs),
	}
}

func (g *objectGraph) SetVertex(vertex vertex, object map[string]any, jsonPath []any) {
	g.ID2vertex[vertex] = vertexData{
		Object:   object,
		JSONPath: jsonPath,
	}
}

func (g *objectGraph) SetArc(source, destination vertex, arcType arcType, attributes arcAttrs) {
	if _, ok := g.ID2vertex[source]; !ok {
		panic(fmt.Sprintf("source vertex %s not found in graph", source))
	}
	arcSet := g.ID2arcSet[source]
	if arcSet == nil {
		arcSet = make(map[arcToVertex]arcAttrs)
		g.ID2arcSet[source] = arcSet
	}
	toV := arcToVertex{
		To:      destination,
		ArcType: arcType,
	}
	arcSet[toV] = attributes
}

func (g *objectGraph) deleteVertexInternal(vertex vertex) {
	if _, ok := g.ID2vertex[vertex]; !ok {
		panic(fmt.Sprintf("vertex %s not found in graph", vertex))
	}
	delete(g.ID2vertex, vertex)
}

func (g *objectGraph) DeleteVertex(vertex vertex) {
	g.deleteVertexInternal(vertex)
	if arcSet, ok := g.ID2arcSet[vertex]; ok {
		panic(fmt.Sprintf("still have arcs, but not expected to have them for %s: %v", vertex, arcSet))
	}
}

func (g *objectGraph) DeleteArc(source, destination vertex, arcType arcType) {
	arcSet := g.ID2arcSet[source]
	toV := arcToVertex{
		To:      destination,
		ArcType: arcType,
	}
	delete(arcSet, toV)
	if len(arcSet) == 0 {
		delete(g.ID2arcSet, source)
	}
}
