/*
Copyright 2026.

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 rbac

import (
	"context"
	"fmt"
	"reflect"
	"strings"

	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"

	infrav1alpha1 "github.com/sandbox0-ai/sandbox0/infra-operator/api/v1alpha1"
	"github.com/sandbox0-ai/sandbox0/infra-operator/internal/controller/pkg/common"
)

const gkeWorkloadIdentityAnnotation = "iam.gke.io/gcp-service-account"

type Reconciler struct {
	Resources *common.ResourceManager
}

func NewReconciler(resources *common.ResourceManager) *Reconciler {
	return &Reconciler{Resources: resources}
}

// ReconcileManagerRBAC reconciles RBAC for the manager service.
func (r *Reconciler) ReconcileManagerRBAC(ctx context.Context, infra *infrav1alpha1.Sandbox0Infra) error {
	name := fmt.Sprintf("%s-manager", infra.Name)
	labels := map[string]string{
		"app.kubernetes.io/name":       "manager",
		"app.kubernetes.io/instance":   infra.Name,
		"app.kubernetes.io/managed-by": "sandbox0infra-operator",
	}

	if err := r.reconcileServiceAccount(ctx, infra, name, labels, nil); err != nil {
		return err
	}

	rules := []rbacv1.PolicyRule{
		{
			APIGroups: []string{"sandbox0.ai"},
			Resources: []string{"sandboxtemplates", "sandboxtemplates/status"},
			Verbs:     []string{"get", "list", "watch", "create", "update", "patch", "delete"},
		},
		{
			APIGroups: []string{""},
			Resources: []string{"pods", "pods/exec", "pods/resize", "pods/status", "services", "serviceaccounts", "configmaps", "secrets", "persistentvolumeclaims", "events", "nodes", "namespaces"},
			Verbs:     []string{"get", "list", "watch", "create", "update", "patch", "delete"},
		},
		{
			APIGroups: []string{""},
			Resources: []string{"pods/log"},
			Verbs:     []string{"get"},
		},
		{
			APIGroups: []string{"apps"},
			Resources: []string{"replicasets", "deployments"},
			Verbs:     []string{"get", "list", "watch", "create", "update", "patch", "delete"},
		},
		{
			APIGroups: []string{"networking.k8s.io"},
			Resources: []string{"networkpolicies"},
			Verbs:     []string{"get", "list", "watch", "create", "update", "patch", "delete"},
		},
		{
			APIGroups: []string{"coordination.k8s.io"},
			Resources: []string{"leases"},
			Verbs:     []string{"get", "list", "watch", "create", "update", "patch", "delete"},
		},
		{
			APIGroups: []string{"admissionregistration.k8s.io"},
			Resources: []string{"validatingwebhookconfigurations", "mutatingwebhookconfigurations"},
			Verbs:     []string{"get", "list", "watch", "create", "update", "patch", "delete"},
		},
	}

	if err := r.reconcileClusterRole(ctx, name, labels, rules); err != nil {
		return err
	}

	return r.reconcileClusterRoleBinding(ctx, infra, name, labels, name, name)
}

// ReconcileNetdRBAC reconciles RBAC for the netd service.
func (r *Reconciler) ReconcileNetdRBAC(ctx context.Context, infra *infrav1alpha1.Sandbox0Infra) error {
	name := fmt.Sprintf("%s-netd", infra.Name)
	labels := map[string]string{
		"app.kubernetes.io/name":       "netd",
		"app.kubernetes.io/instance":   infra.Name,
		"app.kubernetes.io/managed-by": "sandbox0infra-operator",
	}

	if err := r.reconcileServiceAccount(ctx, infra, name, labels, nil); err != nil {
		return err
	}

	rules := []rbacv1.PolicyRule{
		{
			APIGroups: []string{""},
			Resources: []string{"pods", "pods/status", "nodes", "services", "endpoints"},
			Verbs:     []string{"get", "list", "watch", "patch"},
		},
		{
			APIGroups: []string{"discovery.k8s.io"},
			Resources: []string{"endpointslices"},
			Verbs:     []string{"get", "list", "watch"},
		},
	}

	if err := r.reconcileClusterRole(ctx, name, labels, rules); err != nil {
		return err
	}

	return r.reconcileClusterRoleBinding(ctx, infra, name, labels, name, name)
}

// ReconcileSchedulerRBAC reconciles RBAC for the scheduler service.
func (r *Reconciler) ReconcileSchedulerRBAC(ctx context.Context, infra *infrav1alpha1.Sandbox0Infra) error {
	name := fmt.Sprintf("%s-scheduler", infra.Name)
	labels := map[string]string{
		"app.kubernetes.io/name":       "scheduler",
		"app.kubernetes.io/instance":   infra.Name,
		"app.kubernetes.io/managed-by": "sandbox0infra-operator",
	}

	return r.reconcileServiceAccount(ctx, infra, name, labels, nil)
}

// ReconcileStorageProxyRBAC reconciles RBAC for the storage-proxy service.
func (r *Reconciler) ReconcileStorageProxyRBAC(ctx context.Context, infra *infrav1alpha1.Sandbox0Infra) error {
	name := fmt.Sprintf("%s-storage-proxy", infra.Name)
	labels := map[string]string{
		"app.kubernetes.io/name":       "storage-proxy",
		"app.kubernetes.io/instance":   infra.Name,
		"app.kubernetes.io/managed-by": "sandbox0infra-operator",
	}

	if err := r.reconcileServiceAccount(ctx, infra, name, labels, storageWorkloadIdentityAnnotations(infra)); err != nil {
		return err
	}

	rules := []rbacv1.PolicyRule{
		{
			APIGroups: []string{""},
			Resources: []string{"pods", "events"},
			Verbs:     []string{"get", "list", "watch", "create", "patch"},
		},
		{
			APIGroups: []string{""},
			Resources: []string{"nodes"},
			Verbs:     []string{"get"},
		},
	}

	if err := r.reconcileClusterRole(ctx, name, labels, rules); err != nil {
		return err
	}

	return r.reconcileClusterRoleBinding(ctx, infra, name, labels, name, name)
}

// ReconcileCtldRBAC reconciles RBAC for the ctld daemonset.
func (r *Reconciler) ReconcileCtldRBAC(ctx context.Context, infra *infrav1alpha1.Sandbox0Infra) error {
	name := fmt.Sprintf("%s-ctld", infra.Name)
	labels := map[string]string{
		"app.kubernetes.io/name":       "ctld",
		"app.kubernetes.io/instance":   infra.Name,
		"app.kubernetes.io/managed-by": "sandbox0infra-operator",
	}

	if err := r.reconcileServiceAccount(ctx, infra, name, labels, storageWorkloadIdentityAnnotations(infra)); err != nil {
		return err
	}

	rules := []rbacv1.PolicyRule{
		{
			APIGroups: []string{""},
			Resources: []string{"pods"},
			Verbs:     []string{"get", "list", "watch", "update", "patch"},
		},
		{
			APIGroups: []string{""},
			Resources: []string{"pods/resize"},
			Verbs:     []string{"update", "patch"},
		},
	}

	if err := r.reconcileClusterRole(ctx, name, labels, rules); err != nil {
		return err
	}

	return r.reconcileClusterRoleBinding(ctx, infra, name, labels, name, name)
}

// reconcileServiceAccount creates or updates a ServiceAccount for a service.
func (r *Reconciler) reconcileServiceAccount(ctx context.Context, infra *infrav1alpha1.Sandbox0Infra, name string, labels map[string]string, annotations map[string]string) error {
	sa := &corev1.ServiceAccount{
		ObjectMeta: metav1.ObjectMeta{
			Name:        name,
			Namespace:   infra.Namespace,
			Labels:      labels,
			Annotations: annotations,
		},
	}

	if err := ctrl.SetControllerReference(infra, sa, r.Resources.Scheme); err != nil {
		return err
	}

	found := &corev1.ServiceAccount{}
	err := r.Resources.Client.Get(ctx, types.NamespacedName{Name: name, Namespace: infra.Namespace}, found)
	if err != nil && errors.IsNotFound(err) {
		return r.Resources.Client.Create(ctx, sa)
	} else if err != nil {
		return err
	}

	updated := found.DeepCopy()
	updated.Labels = labels
	if annotations != nil {
		updated.Annotations = reconcileServiceAccountAnnotations(updated.Annotations, annotations)
	}
	if reflect.DeepEqual(found.Labels, updated.Labels) && reflect.DeepEqual(found.Annotations, updated.Annotations) {
		return nil
	}
	return r.Resources.Client.Update(ctx, updated)
}

func storageWorkloadIdentityAnnotations(infra *infrav1alpha1.Sandbox0Infra) map[string]string {
	if infra == nil || infra.Spec.Storage == nil || infra.Spec.Storage.Type != infrav1alpha1.StorageTypeGCS || infra.Spec.Storage.GCS == nil {
		return map[string]string{}
	}
	email := strings.TrimSpace(infra.Spec.Storage.GCS.WorkloadIdentityServiceAccountEmail)
	if email == "" {
		return map[string]string{}
	}
	return map[string]string{gkeWorkloadIdentityAnnotation: email}
}

func reconcileServiceAccountAnnotations(current, desired map[string]string) map[string]string {
	next := common.CloneStringMap(current)
	if next == nil {
		next = map[string]string{}
	}
	if value := strings.TrimSpace(desired[gkeWorkloadIdentityAnnotation]); value != "" {
		next[gkeWorkloadIdentityAnnotation] = value
	} else {
		delete(next, gkeWorkloadIdentityAnnotation)
	}
	if len(next) == 0 {
		return nil
	}
	return next
}

// reconcileClusterRole creates or updates a ClusterRole for a service.
func (r *Reconciler) reconcileClusterRole(ctx context.Context, name string, labels map[string]string, rules []rbacv1.PolicyRule) error {
	role := &rbacv1.ClusterRole{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: labels,
		},
		Rules: rules,
	}

	found := &rbacv1.ClusterRole{}
	err := r.Resources.Client.Get(ctx, types.NamespacedName{Name: name}, found)
	if err != nil && errors.IsNotFound(err) {
		return r.Resources.Client.Create(ctx, role)
	} else if err != nil {
		return err
	}

	return r.Resources.UpdateObjectIfChanged(ctx, found, func() {
		found.Labels = labels
		found.Rules = rules
	})
}

// reconcileClusterRoleBinding creates or updates a ClusterRoleBinding for a service.
func (r *Reconciler) reconcileClusterRoleBinding(ctx context.Context, infra *infrav1alpha1.Sandbox0Infra, name string, labels map[string]string, roleName string, saName string) error {
	binding := &rbacv1.ClusterRoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: labels,
		},
		RoleRef: rbacv1.RoleRef{
			APIGroup: "rbac.authorization.k8s.io",
			Kind:     "ClusterRole",
			Name:     roleName,
		},
		Subjects: []rbacv1.Subject{
			{
				Kind:      "ServiceAccount",
				Name:      saName,
				Namespace: infra.Namespace,
			},
		},
	}

	found := &rbacv1.ClusterRoleBinding{}
	err := r.Resources.Client.Get(ctx, types.NamespacedName{Name: name}, found)
	if err != nil && errors.IsNotFound(err) {
		return r.Resources.Client.Create(ctx, binding)
	} else if err != nil {
		return err
	}

	return r.Resources.UpdateObjectIfChanged(ctx, found, func() {
		found.Labels = labels
		found.RoleRef = binding.RoleRef
		found.Subjects = binding.Subjects
	})
}
