/*
Copyright 2024 The Kubernetes Authors.

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 generic

import (
	"context"
	"errors"
	"fmt"

	admissionregistrationv1 "k8s.io/api/admissionregistration/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/runtime/schema"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apiserver/pkg/admission"
	"k8s.io/apiserver/pkg/admission/initializer"
	"k8s.io/apiserver/pkg/admission/plugin/policy/matching"
	"k8s.io/apiserver/pkg/authorization/authorizer"
	"k8s.io/apiserver/pkg/features"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
)

// H is the Hook type generated by the source and consumed by the dispatcher.
// !TODO: Just pass in a Plugin[H] with accessors to all this information
type apiSourceFactory[H any] func(informers.SharedInformerFactory, kubernetes.Interface, dynamic.Interface, meta.RESTMapper) Source[H]
type dispatcherFactory[H any] func(authorizer.Authorizer, *matching.Matcher, kubernetes.Interface) Dispatcher[H]

// ReloadableSource extends HookSource with a method to run a reload loop
// that watches for configuration changes and blocks until the context is canceled.
type ReloadableSource[H any] interface {
	HookSource[H]
	// RunReloadLoop watches for configuration changes and reloads when detected.
	// It blocks until ctx is canceled.
	RunReloadLoop(ctx context.Context)
}

// StaticSourceFactory creates a static policy source from a manifest directory.
// The returned Source should have LoadInitial() already called.
type StaticSourceFactory[H any] func(manifestsDir string) (ReloadableSource[H], error)

// admissionResources is the list of resources related to CEL-based admission
// features.
var admissionResources = []schema.GroupResource{
	{Group: admissionregistrationv1.GroupName, Resource: "validatingadmissionpolicies"},
	{Group: admissionregistrationv1.GroupName, Resource: "validatingadmissionpolicybindings"},
	{Group: admissionregistrationv1.GroupName, Resource: "mutatingadmissionpolicies"},
	{Group: admissionregistrationv1.GroupName, Resource: "mutatingadmissionpolicybindings"},
}

// AdmissionPolicyManager is an abstract admission plugin with all the
// infrastructure to define Admit or Validate on-top.
type Plugin[H any] struct {
	*admission.Handler

	apiSourceFactory    apiSourceFactory[H]
	dispatcherFactory   dispatcherFactory[H]
	staticSourceFactory StaticSourceFactory[H]

	source     Source[H]
	dispatcher Dispatcher[H]
	matcher    *matching.Matcher

	// staticSource holds a reference to only the static (manifest-based) policy source.
	// This is used in Dispatch to run static hooks for admission config resources
	// (e.g., VAP/MAP/VAPB/MAPB), allowing static policies to protect those resources
	// without the circular dependency concern that applies to API-based policies.
	staticSource HookSource[H]

	// staticManifestsDir is the path to a directory containing static policy manifests.
	// When set and the feature gate is enabled, policies are also loaded from this directory.
	staticManifestsDir string

	// apiServerID is the identity of this API server instance, used for metrics labeling.
	apiServerID string

	informerFactory informers.SharedInformerFactory
	client          kubernetes.Interface
	restMapper      meta.RESTMapper
	dynamicClient   dynamic.Interface
	// admissionConfigResources are admission configuration resources (VAP/MAP/VAPB/MAPB).
	// API-based policies skip these to prevent circular dependencies, but static policies
	// are safe to evaluate on them.
	admissionConfigResources sets.Set[schema.GroupResource]
	// excludedResources are non-persisted resources (auth/authz reviews) that must be
	// excluded from ALL policy evaluation, including static policies.
	excludedResources sets.Set[schema.GroupResource]
	stopCh            <-chan struct{}
	authorizer        authorizer.Authorizer
	enabled           bool
}

var (
	_ initializer.WantsExternalKubeInformerFactory = &Plugin[any]{}
	_ initializer.WantsExternalKubeClientSet       = &Plugin[any]{}
	_ initializer.WantsRESTMapper                  = &Plugin[any]{}
	_ initializer.WantsDynamicClient               = &Plugin[any]{}
	_ initializer.WantsDrainedNotification         = &Plugin[any]{}
	_ initializer.WantsAPIServerID                 = &Plugin[any]{}
	_ initializer.WantsAuthorizer                  = &Plugin[any]{}
	_ initializer.WantsExcludedAdmissionResources  = &Plugin[any]{}
	_ admission.InitializationValidator            = &Plugin[any]{}
)

func NewPlugin[H any](
	handler *admission.Handler,
	apiSourceFactory apiSourceFactory[H],
	dispatcherFactory dispatcherFactory[H],
) *Plugin[H] {
	return &Plugin[H]{
		Handler:           handler,
		apiSourceFactory:  apiSourceFactory,
		dispatcherFactory: dispatcherFactory,

		// admission config resources are skipped by API-based policies (circular dep)
		// but evaluated by static policies
		admissionConfigResources: sets.New(admissionResources...),
		excludedResources:        sets.New[schema.GroupResource](),
	}
}

func (c *Plugin[H]) SetExternalKubeInformerFactory(f informers.SharedInformerFactory) {
	c.informerFactory = f
}

func (c *Plugin[H]) SetExternalKubeClientSet(client kubernetes.Interface) {
	c.client = client
}

func (c *Plugin[H]) SetRESTMapper(mapper meta.RESTMapper) {
	c.restMapper = mapper
}

func (c *Plugin[H]) SetDynamicClient(client dynamic.Interface) {
	c.dynamicClient = client
}

func (c *Plugin[H]) SetDrainedNotification(stopCh <-chan struct{}) {
	c.stopCh = stopCh
}

func (c *Plugin[H]) SetAuthorizer(authorizer authorizer.Authorizer) {
	c.authorizer = authorizer
}

func (c *Plugin[H]) SetMatcher(matcher *matching.Matcher) {
	c.matcher = matcher
}

func (c *Plugin[H]) SetEnabled(enabled bool) {
	c.enabled = enabled
}

func (c *Plugin[H]) SetExcludedAdmissionResources(excludedResources []schema.GroupResource) {
	c.excludedResources.Insert(excludedResources...)
}

// SetStaticManifestsDir sets the directory containing static policy manifests.
func (c *Plugin[H]) SetStaticManifestsDir(dir string) {
	c.staticManifestsDir = dir
}

// GetStaticManifestsDir returns the directory containing static policy manifests.
func (c *Plugin[H]) GetStaticManifestsDir() string {
	return c.staticManifestsDir
}

// SetStaticSourceFactory sets the factory for creating static policy sources.
// This should be called before ValidateInitialization.
func (c *Plugin[H]) SetStaticSourceFactory(factory StaticSourceFactory[H]) {
	c.staticSourceFactory = factory
}

// SetAPIServerID implements the WantsAPIServerID interface.
// The API server ID is used for metrics labeling and must be set before
// ValidateInitialization is called.
func (c *Plugin[H]) SetAPIServerID(id string) {
	c.apiServerID = id
}

// GetAPIServerID returns the stored API server ID.
func (c *Plugin[H]) GetAPIServerID() string {
	return c.apiServerID
}

// ValidateInitialization - once clientset and informer factory are provided, creates and starts the admission controller
func (c *Plugin[H]) ValidateInitialization() error {
	// By default enabled is set to false. It is up to types which embed this
	// struct to set it to true (if feature gate is enabled, or other conditions)
	if !c.enabled {
		return nil
	}
	if c.Handler == nil {
		return errors.New("missing handler")
	}
	if c.informerFactory == nil {
		return errors.New("missing informer factory")
	}
	if c.client == nil {
		return errors.New("missing kubernetes client")
	}
	if c.restMapper == nil {
		return errors.New("missing rest mapper")
	}
	if c.dynamicClient == nil {
		return errors.New("missing dynamic client")
	}
	if c.stopCh == nil {
		return errors.New("missing stop channel")
	}
	if c.authorizer == nil {
		return errors.New("missing authorizer")
	}

	// Guard: if static manifests dir is set but feature gate is off, return an error
	if len(c.staticManifestsDir) > 0 && !utilfeature.DefaultFeatureGate.Enabled(features.ManifestBasedAdmissionControlConfig) {
		return fmt.Errorf("static policy manifests dir %q configured but %s feature gate is not enabled", c.staticManifestsDir, features.ManifestBasedAdmissionControlConfig)
	}

	namespaceInformer := c.informerFactory.Core().V1().Namespaces()

	// Construct matcher once, avoiding overwriting if already set.
	if c.matcher == nil {
		c.matcher = matching.NewMatcher(namespaceInformer.Lister(), c.client)
	}
	if err := c.matcher.ValidateInitialization(); err != nil {
		return err
	}

	// Create a shared context tied to server shutdown for all background goroutines.
	pluginContext, pluginContextCancel := context.WithCancel(context.Background())
	go func() {
		defer pluginContextCancel()
		<-c.stopCh
	}()

	// Construct source once, avoiding overwriting if already set.
	if c.source == nil {
		apiSource := c.apiSourceFactory(c.informerFactory, c.client, c.dynamicClient, c.restMapper)

		if len(c.staticManifestsDir) > 0 {
			if c.staticSourceFactory == nil {
				return fmt.Errorf("static policy manifests configured in %q but no static source factory is set", c.staticManifestsDir)
			}
			staticSource, err := c.staticSourceFactory(c.staticManifestsDir)
			if err != nil {
				return fmt.Errorf("failed to load static policy manifests from %q: %w", c.staticManifestsDir, err)
			}
			// Use composite source that combines static + API sources
			c.source = NewCompositePolicySource(staticSource, apiSource)
			// Keep a reference to the static source so Dispatch can run static hooks
			// for resources that are normally excluded from policy evaluation.
			c.staticSource = staticSource
			// Start the file watcher in a background goroutine
			go staticSource.RunReloadLoop(pluginContext)
		} else {
			// Use only API-based source
			c.source = apiSource
		}

		go func() {
			err := c.source.Run(pluginContext)
			if err != nil && !errors.Is(err, context.Canceled) {
				utilruntime.HandleError(fmt.Errorf("policy source context unexpectedly closed: %w", err))
			}
		}()

		c.SetReadyFunc(func() bool {
			return namespaceInformer.Informer().HasSynced() && c.source.HasSynced()
		})
	}

	// Construct dispatcher once, avoiding overwriting if already set.
	if c.dispatcher == nil {
		c.dispatcher = c.dispatcherFactory(c.authorizer, c.matcher, c.client)

		err := c.dispatcher.Start(pluginContext)
		if err != nil && !errors.Is(err, context.Canceled) {
			utilruntime.HandleError(fmt.Errorf("policy dispatcher context unexpectedly closed: %w", err))
		}
	}

	return nil
}

func (c *Plugin[H]) Dispatch(
	ctx context.Context,
	a admission.Attributes,
	o admission.ObjectInterfaces,
) (err error) {
	if !c.enabled {
		return nil
	}

	gr := a.GetResource().GroupResource()

	if c.isExcludedFromAllHooks(gr) {
		// Non-persisted resources (auth/authz reviews) are excluded from all policy
		// evaluation to avoid breaking the cluster.
		return nil
	} else if c.isExcludedFromAPIHooks(gr) {
		// Admission config resources (VAP/MAP/VAPB/MAPB) are excluded from API-based
		// policies to prevent circular dependencies. However, static (manifest-based)
		// policies are safe to evaluate since they don't have self-referential concerns.
		if c.staticSource != nil {
			if !c.staticSource.HasSynced() {
				return admission.NewForbidden(a, fmt.Errorf("not yet ready to handle request"))
			}
			return c.dispatcher.Dispatch(ctx, a, o, c.staticSource.Hooks())
		}
		return nil
	}

	if !c.WaitForReady() {
		return admission.NewForbidden(a, fmt.Errorf("not yet ready to handle request"))
	}

	return c.dispatcher.Dispatch(ctx, a, o, c.source.Hooks())
}

// isExcludedFromAllHooks returns true for non-persisted resources (auth/authz reviews)
// that must not be intercepted by any policy hook.
func (c *Plugin[H]) isExcludedFromAllHooks(gr schema.GroupResource) bool {
	return c.excludedResources.Has(gr)
}

// isExcludedFromAPIHooks returns true for admission config resources (VAP/MAP/VAPB/MAPB)
// that are excluded from API-based hooks but may be evaluated by static hooks.
func (c *Plugin[H]) isExcludedFromAPIHooks(gr schema.GroupResource) bool {
	return c.admissionConfigResources.Has(gr)
}
