package jumper

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	jmpent "github.com/tidbcloud/devops/entity/jump"
	"github.com/tidbcloud/devops/pkg/infra/log"
	"github.com/tidbcloud/devops/pkg/infra/utils/container"
	"github.com/tidbcloud/devops/pkg/service/jump/jumppod/types"
	rbacv1 "k8s.io/api/rbac/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	ctrlutil "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
	sigyaml "sigs.k8s.io/yaml"
)

const (
	JumpReadOnlyAggregationLabelPrefix = "jump-service.devops.cloud.pingcap.com/aggregate-to-"
)

type JumpRoleCreator interface {
	// PrepareRoles creates or updates shared cluster roles for jump.
	PrepareRoles(ctx context.Context, ctrl client.Client) error
}

type jumpRoleCreator struct {
	dir string

	AggregateRoles []*rbacv1.ClusterRole
	CoreRoles      []*rbacv1.ClusterRole
}

func NewJumpRoleCreator(dir string) (*jumpRoleCreator, error) {
	b := &jumpRoleCreator{
		dir: dir,
	}
	if err := b.Load(); err != nil {
		return nil, err
	}
	return b, nil
}

func (c *jumpRoleCreator) PrepareRoles(ctx context.Context, ctrl client.Client) error {
	commonLabels := types.NewLabels(nil)
	commonAnnos := types.NewAnnotations(nil)

	// create or update core roles
	for _, oriCoreRole := range c.CoreRoles {
		expCoreRole := oriCoreRole.DeepCopy()
		coreRole := &rbacv1.ClusterRole{
			ObjectMeta: metav1.ObjectMeta{
				Name:      expCoreRole.Name,
				Namespace: expCoreRole.Namespace,
			},
		}
		_, err := ctrlutil.CreateOrUpdate(ctx, ctrl, coreRole, func() error {
			if coreRole.Labels == nil {
				coreRole.Labels = map[string]string{}
			}
			commonLabels.InjectTo(coreRole.Labels)

			if coreRole.Annotations == nil {
				coreRole.Annotations = map[string]string{}
			}
			commonAnnos.InjectTo(coreRole.Annotations)

			coreRole.AggregationRule = expCoreRole.AggregationRule
			return nil
		})
		if err != nil {
			return fmt.Errorf("create or update cluster role %s failed: %w", coreRole.Name, err)
		}
	}

	// create or update aggregate roles
	seen := container.NewSet[string]()
	for _, orgAggRole := range c.AggregateRoles {
		seen.Insert(orgAggRole.Name)

		expAggRule := orgAggRole.DeepCopy()
		aggRole := &rbacv1.ClusterRole{
			ObjectMeta: metav1.ObjectMeta{
				Name:      expAggRule.Name,
				Namespace: expAggRule.Namespace,
			},
		}

		_, err := ctrlutil.CreateOrUpdate(ctx, ctrl, aggRole, func() error {
			if aggRole.Labels == nil {
				aggRole.Labels = map[string]string{}
			}
			commonLabels.InjectTo(aggRole.Labels)
			for k, v := range expAggRule.Labels { // inject original labels
				aggRole.Labels[k] = v
			}

			if aggRole.Annotations == nil {
				aggRole.Annotations = map[string]string{}
			}
			commonAnnos.InjectTo(aggRole.Annotations)

			aggRole.Rules = expAggRule.Rules
			return nil
		})
		if err != nil {
			return fmt.Errorf("create or update cluster role %s failed: %w", aggRole.Name, err)
		}
	}

	// delete stale aggregate roles
	roleList := &rbacv1.ClusterRoleList{}
	if err := ctrl.List(ctx, roleList, client.MatchingLabels{
		types.AppLabelKey: types.AppLabelValue,
	}); err != nil {
		return fmt.Errorf("list cluster roles failed: %w", err)
	}
	for i := range roleList.Items {
		role := &roleList.Items[i]

		if seen.Has(role.Name) {
			continue
		}

		isAggRole := false
		for k := range role.Labels {
			if strings.HasPrefix(k, JumpReadOnlyAggregationLabelPrefix) {
				isAggRole = true
				break
			}
		}
		if !isAggRole {
			continue
		}

		if err := ctrl.Delete(ctx, role); err != nil {
			log.FromContext(ctx).Error(err, "delete stale aggregate cluster role failed", "role", role.Name)
		}
	}

	return nil
}

func (b *jumpRoleCreator) JumpAggregateLabel(roleName string) string {
	return JumpReadOnlyAggregationLabelPrefix + roleName
}

func (b *jumpRoleCreator) Load() error {
	// build core rules
	for _, priv := range jmpent.JumpPrivilegeValues {
		coreRole := &rbacv1.ClusterRole{
			ObjectMeta: metav1.ObjectMeta{
				Name:      JumpRoleName(priv),
				Namespace: "",
			},
			Rules: []rbacv1.PolicyRule{},
			AggregationRule: &rbacv1.AggregationRule{
				ClusterRoleSelectors: []metav1.LabelSelector{
					{
						MatchLabels: map[string]string{
							b.JumpAggregateLabel(JumpRoleName(priv)): "true",
						},
					},
				},
			},
		}
		coreGWRole := &rbacv1.ClusterRole{
			ObjectMeta: metav1.ObjectMeta{
				Name:      JumpGatewayRoleName(priv),
				Namespace: "",
			},
			Rules: []rbacv1.PolicyRule{},
			AggregationRule: &rbacv1.AggregationRule{
				ClusterRoleSelectors: []metav1.LabelSelector{
					{
						MatchLabels: map[string]string{
							b.JumpAggregateLabel(JumpGatewayRoleName(priv)): "true",
						},
					},
				},
			},
		}

		b.CoreRoles = append(b.CoreRoles, coreRole, coreGWRole)
	}

	// load aggregate rules

	entries, err := os.ReadDir(b.dir)
	if err != nil {
		return fmt.Errorf("read roles dir %s failed: %w", b.dir, err)
	}
	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Name() < entries[j].Name()
	})

	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		ext := strings.ToLower(filepath.Ext(entry.Name()))
		if ext != ".yaml" && ext != ".yml" {
			continue
		}

		path := filepath.Join(b.dir, entry.Name())
		data, err := os.ReadFile(path)
		if err != nil {
			return fmt.Errorf("read role file %s failed: %w", path, err)
		}

		role := &rbacv1.ClusterRole{}
		if err := sigyaml.Unmarshal(data, role); err != nil {
			return fmt.Errorf("decode clusterrole from %s failed: %w", path, err)
		}

		b.AggregateRoles = append(b.AggregateRoles, role)
	}

	return nil
}