package test

import (
	"context"
	"errors"

	"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"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/apiutil"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

// ErrorInjectingFakeClient is a fake client that can be used to inject errors into the K8S client
type ErrorInjectingFakeClient struct {
	client.WithWatch
	withSubResourceSimulation bool
	CustomSubResourceWriter   *CustomSubResourceWriter
	withObjects               []client.Object
	withStatusSubresource     []client.Object
	failOnCreate              bool
	failOnUpdate              bool
	failOnList                bool
	failOnDelete              bool
	failOnGet                 bool
	failOnStatus              bool
}

func (c *ErrorInjectingFakeClient) Status() client.SubResourceWriter {
	if c.failOnStatus {
		// use the custom subresource writer to inject errors
		if c.CustomSubResourceWriter == nil {
			return c.NewCustomSubResourceWriter()
		}
		return c.CustomSubResourceWriter
	}
	return c.WithWatch.Status()
}

func NewErrorInjectingFakeClient(scheme *runtime.Scheme, withSubResourceSimulation bool, initObjects ...client.Object) *ErrorInjectingFakeClient {
	builder := fake.NewClientBuilder().
		WithScheme(scheme).
		WithObjects(initObjects...).
		WithStatusSubresource(initObjects...)
	return &ErrorInjectingFakeClient{
		withSubResourceSimulation: withSubResourceSimulation,
		WithWatch:                 builder.Build(),
		withObjects:               initObjects,
		withStatusSubresource:     initObjects,
	}
}

// Copied from sigs.k8s.io/controller-runtime@v0.16.3/pkg/client/fake/client.go#1217
var inTreeResourcesWithStatus = []schema.GroupVersionKind{
	{Version: "v1", Kind: "Namespace"},
	{Version: "v1", Kind: "Node"},
	{Version: "v1", Kind: "PersistentVolumeClaim"},
	{Version: "v1", Kind: "PersistentVolume"},
	{Version: "v1", Kind: "Pod"},
	{Version: "v1", Kind: "ReplicationController"},
	{Version: "v1", Kind: "Service"},

	{Group: "apps", Version: "v1", Kind: "Deployment"},
	{Group: "apps", Version: "v1", Kind: "DaemonSet"},
	{Group: "apps", Version: "v1", Kind: "ReplicaSet"},
	{Group: "apps", Version: "v1", Kind: "StatefulSet"},

	{Group: "autoscaling", Version: "v1", Kind: "HorizontalPodAutoscaler"},

	{Group: "batch", Version: "v1", Kind: "CronJob"},
	{Group: "batch", Version: "v1", Kind: "Job"},

	{Group: "certificates.k8s.io", Version: "v1", Kind: "CertificateSigningRequest"},

	{Group: "networking.k8s.io", Version: "v1", Kind: "Ingress"},
	{Group: "networking.k8s.io", Version: "v1", Kind: "NetworkPolicy"},

	{Group: "policy", Version: "v1", Kind: "PodDisruptionBudget"},

	{Group: "storage.k8s.io", Version: "v1", Kind: "VolumeAttachment"},

	{Group: "apiextensions.k8s.io", Version: "v1", Kind: "CustomResourceDefinition"},

	{Group: "flowcontrol.apiserver.k8s.io", Version: "v1beta2", Kind: "FlowSchema"},
	{Group: "flowcontrol.apiserver.k8s.io", Version: "v1beta2", Kind: "PriorityLevelConfiguration"},
}

func (c *ErrorInjectingFakeClient) Create(ctx context.Context, obj client.Object, opts ...client.CreateOption) error {
	if c.failOnCreate {
		return errors.New("create error")
	}

	if err := c.WithWatch.Create(ctx, obj, opts...); err != nil {
		return err
	}
	if !c.withSubResourceSimulation {
		return nil
	}

	c.withObjects = append(c.withObjects, obj)
	err := c.recreateFake(ctx)
	if err != nil {
		return err
	}

	return nil
}

func (c *ErrorInjectingFakeClient) Delete(ctx context.Context, obj client.Object, opts ...client.DeleteOption) error {
	if c.failOnDelete {
		return errors.New("delete error")
	}

	if err := c.WithWatch.Delete(ctx, obj, opts...); err != nil {
		return err
	}

	return nil
}

func (c *ErrorInjectingFakeClient) Update(ctx context.Context, obj client.Object, opts ...client.UpdateOption) error {
	if c.failOnUpdate {
		return errors.New("update error")
	}
	return c.WithWatch.Update(ctx, obj, opts...)
}

// create func for List method
func (c *ErrorInjectingFakeClient) List(ctx context.Context, list client.ObjectList, opts ...client.ListOption) error {
	if c.failOnList {
		return errors.New("unsupported list type")
	}
	return c.WithWatch.List(ctx, list, opts...)
}

// create func for Get method
func (c *ErrorInjectingFakeClient) Get(ctx context.Context, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error {
	if c.failOnGet {
		return errors.New("unsupported get type")
	}
	return c.WithWatch.Get(ctx, key, obj, opts...)
}

func (c *ErrorInjectingFakeClient) recreateFake(ctx context.Context) error {
	// This is a pretty disgusting hack to get around the fact that the fake client doesn't support creating status subresources
	// Discussed here https://github.com/kubernetes-sigs/controller-runtime/issues/2386#issuecomment-1607768830
	// and here https://github.com/kubernetes-sigs/controller-runtime/issues/2362#issuecomment-1699415588
	// and here https://stackoverflow.com/questions/77489441/go-k8s-controller-runtime-upgrade-fake-client-lacks-functionality

	// We need to collect all the objects from the fake client, and then re-create the fake client with the status subresource for the created object

	gvks := sets.New(inTreeResourcesWithStatus...)
	for _, o := range c.withObjects {
		gvk, err := apiutil.GVKForObject(o, c.Scheme())
		if err != nil {
			return err
		}
		gvks.Insert(gvk)
	}

	var objs []client.Object
	for _, gvk := range gvks.UnsortedList() {
		objList := &unstructured.UnstructuredList{}
		objList.SetGroupVersionKind(schema.GroupVersionKind{Group: gvk.Group, Version: gvk.Version, Kind: gvk.Kind})

		err := c.List(ctx, objList)
		if err != nil {
			return err
		}

		for _, o := range objList.Items {
			objs = append(objs, o.DeepCopy())
		}
	}

	initObjs := sets.New(objs...).UnsortedList()
	c.withObjects = initObjs
	c.withStatusSubresource = initObjs

	c.WithWatch = fake.NewClientBuilder().
		WithScheme(c.Scheme()).
		WithObjects(c.withObjects...).
		WithStatusSubresource(c.withStatusSubresource...).
		Build()
	return nil
}

// SetFailOnCreate sets the flag to fail on create
func (c *ErrorInjectingFakeClient) SetFailOnCreate() {
	c.failOnCreate = true
}

// SetFailOnUpdate sets the flag to fail on update
func (c *ErrorInjectingFakeClient) SetFailOnUpdate() {
	c.failOnUpdate = true
}

// SetFailOnList sets the flag to fail on list
func (c *ErrorInjectingFakeClient) SetFailOnList() {
	c.failOnList = true
}

// SetFailOnDelete sets the flag to fail on delete
func (c *ErrorInjectingFakeClient) SetFailOnDelete() {
	c.failOnDelete = true
}

// SetFailOnGet sets the flag to fail on get
func (c *ErrorInjectingFakeClient) SetFailOnGet() {
	c.failOnGet = true
}

// SetFailOnStatus sets the flag to fail on status
func (c *ErrorInjectingFakeClient) SetFailOnStatus() {
	c.failOnStatus = true
}

// CustomSubResourceWriter is a custom implementation of the client.SubResourceWriter interface
type CustomSubResourceWriter struct {
	client       client.Client
	failOnUpdate bool
	failOnCreate bool
	failOnPatch  bool
}

// setFailOnUpdate sets the flag to fail on update
func (w *CustomSubResourceWriter) SetFailOnUpdate() {
	w.failOnUpdate = true
}

// setFailOnCreate sets the flag to fail on create
func (w *CustomSubResourceWriter) SetFailOnCreate() {
	w.failOnCreate = true
}

// setFailOnPatch sets the flag to fail on patch
func (w *CustomSubResourceWriter) SetFailOnPatch() {
	w.failOnPatch = true
}

func (w *CustomSubResourceWriter) Create(ctx context.Context, obj client.Object, subResource client.Object, opts ...client.SubResourceCreateOption) error {
	if w.failOnCreate {
		return errors.New("create error")
	}
	return w.client.SubResource(obj.GetResourceVersion()).Create(ctx, obj, subResource, opts...)
}

func (w *CustomSubResourceWriter) Update(ctx context.Context, obj client.Object, opts ...client.SubResourceUpdateOption) error {
	if w.failOnUpdate {
		return errors.New("update error")
	}
	return w.client.SubResource(obj.GetResourceVersion()).Update(ctx, obj, opts...)
}

func (w *CustomSubResourceWriter) Patch(ctx context.Context, obj client.Object, patch client.Patch, opts ...client.SubResourcePatchOption) error {
	if w.failOnPatch {
		return errors.New("patch error")
	}
	return w.client.SubResource(obj.GetResourceVersion()).Patch(ctx, obj, patch, opts...)
}

func (c *ErrorInjectingFakeClient) NewCustomSubResourceWriter() *CustomSubResourceWriter {
	c.CustomSubResourceWriter = &CustomSubResourceWriter{
		client: c.WithWatch,
	}
	return c.CustomSubResourceWriter
}