package server

import (
	"context"
	"fmt"
	"strings"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow/dag"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"github.com/tidbcloud/cluster-service-ng/pkg/tioperator"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
)

func (s *serverImpl) SetAuditLogSetting(ctx context.Context, tidbClusterId string, auditLogSetting *entity.AuditLogSetting) (*entity.AuditLogSetting, error) {
	if tidbClusterId == "" {
		return nil, status.Errorf(codes.InvalidArgument, "tidb cluster id is required")
	}

	if auditLogSetting == nil {
		return nil, status.Errorf(codes.InvalidArgument, "audit log setting is required")
	}

	// check audit log setting
	switch strings.ToUpper(s.cloudProvider) {
	case apiv1.ProviderType_AWS.String():
		if err := utils.ValidateAWSRoleARN(auditLogSetting.GetProviderIdentity().GetCloudRoleIdentity()); err != nil {
			return nil, status.Errorf(codes.InvalidArgument, "invalid aws role arn: %s", err)
		}
	case apiv1.ProviderType_GCP.String():
		return nil, status.Errorf(codes.Unimplemented, "GCP cloud provider is not supported yet")
	case apiv1.ProviderType_AZURE.String():
		// Azure supports two auth modes:
		// 1. non-empty cloud_role_identity => SAS token auth
		// 2. empty cloud_role_identity => kubelet identity auth
	case apiv1.ProviderType_ALICLOUD.String():
		if err := utils.ValidateAliCloudRoleARN(auditLogSetting.GetProviderIdentity().GetCloudRoleIdentity()); err != nil {
			return nil, status.Errorf(codes.InvalidArgument, "invalid alicloud role arn: %s", err)
		}
	default:
		return nil, status.Errorf(codes.Unimplemented, "cloud provider %s is not supported", s.cloudProvider)
	}

	var dagExecId string
	var dagSubmitted bool
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		// list tidb groups
		var err error
		tidbCluster, err := s.tidbClusterRepo.Load(ctx, tidbClusterId, false)
		if err != nil {
			if status.Code(err) == codes.NotFound {
				return status.Errorf(codes.FailedPrecondition, "tidb cluster %s not found", tidbClusterId)
			}
			return status.Errorf(codes.Internal, "failed to load tidb cluster: %s", err)
		}

		if tidbCluster.GetPhase() == apiv1.Cluster_CREATING {
			return status.Errorf(codes.FailedPrecondition, "tidb cluster %s is creating, please try again later", tidbCluster.GetId())
		}

		tidbCluster.AuditLogSetting = auditLogSetting

		tidbGroups, err := s.tidbGroupRepo.List(ctx, nil, nil, []string{tidbCluster.GetId()}, nil, false)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to list tidb groups: %s", err)
		}
		if len(tidbGroups) == 0 {
			return status.Errorf(codes.FailedPrecondition, "tidb cluster %s has no associated tidb group", tidbCluster.GetId())
		}

		// check if the sharedPool is locked
		unlockFunc, err := s.lock(ctx, tidbCluster.GetSharedPoolId())
		if err != nil {
			tilog.Error(ctx, "failed to lock sharedPool", zap.String("sharedPoolId", tidbCluster.GetSharedPoolId()), zap.Error(err))
			return err
		}
		defer func() {
			if err := unlockFunc(); err != nil {
				tilog.Error(ctx, "failed to unlock sharedPool", zap.String("sharedPoolId", tidbCluster.GetSharedPoolId()), zap.Error(err))
			}
		}()

		// get shared first
		sharedPool, err := s.sharedPoolRepo.Load(ctx, tidbGroups[0].GetSharedPoolId(), false)
		if err != nil {
			if status.Code(err) != codes.NotFound {
				return status.Errorf(codes.Internal, "failed to load sharedPool: %s", err)
			}
			// not found, check if it's in creating
			return status.Errorf(codes.NotFound, "sharedPool %s not found", tidbGroups[0].GetSharedPoolId())
		}
		if err := guardSharedPoolMutation(sharedPool, sharedPoolMutationBlockedWhilePaused); err != nil {
			return err
		}

		// merge sharedPool status from dag task
		err = s.mergeSharedPoolStatusFromDagTask(ctx, sharedPool)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to merge sharedPool status from dag task: %s", err)
		}

		updateMap := make(map[string]string)

		tidbGroupIds := make([]string, 0, len(tidbGroups))
		for _, tidbGroup := range tidbGroups {
			// get running tidb group from dag
			dagIns, err := s.getRunningDagIns(ctx, []string{
				dag.PlanCreateTiDBGroup,
				dag.PlanUpdateTiDBGroup,
				dag.PlanDeleteTiDBGroup,
			}, []string{tidbGroup.GetSharedPoolId()}, nil, nil, []string{tidbGroup.GetId()})
			if err != nil {
				return status.Errorf(codes.Internal, "failed to get running tidb group dag instance: %s", err)
			}

			for _, dagIns := range dagIns {
				switch dagIns.ExecTask.RootName {
				case dag.PlanCreateTiDBGroup:
					return status.Errorf(codes.FailedPrecondition, "tidb group %s is creating, please try again later", tidbGroup.GetId())
				case dag.PlanUpdateTiDBGroup:
					return status.Errorf(codes.FailedPrecondition, "tidb group %s is updating, please try again later", tidbGroup.GetId())
				case dag.PlanDeleteTiDBGroup:
					return status.Errorf(codes.FailedPrecondition, "tidb group %s is deleting", tidbGroup.GetId())
				}
			}

			// set tidb comp state modifying
			for _, comp := range sharedPool.GetComponentGroups() {
				if comp.GetId() == tidbGroup.GetId() && comp.GetComponentType() == apiv1.ComponentType_TIDB {
					comp.Phase = apiv1.Component_MODIFYING
					break
				}
			}
			tidbGroupIds = append(tidbGroupIds, tidbGroup.GetId())

			// set update option
			for _, comp := range sharedPool.GetComponentGroups() {
				if comp.GetId() == tidbGroup.GetId() && comp.GetComponentType() == apiv1.ComponentType_TIDB {
					updateMap[entity.GetUpdateComponentKey(comp)] = string(tioperator.OpsUpdateTiDBAuditLog)
					break
				}
			}
		}

		sharedPoolStr, err := protojson.Marshal(sharedPool)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to marshal shard pool: %s", err)
		}

		// save tidb cluster
		err = s.tidbClusterRepo.Save(ctx, tidbCluster)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to save tidb cluster: %s", err)
		}

		if len(tidbGroupIds) == 0 {
			tilog.Warn(ctx, "tidb cluster has no tidb group to update audit log setting", zap.String("tidbClusterId", tidbCluster.GetId()))
			return nil
		}
		args := actions.ActionArgs{
			workflow.SharedPoolIdKey:   actions.ActionArgString(sharedPool.GetId()),
			workflow.SharedPoolKey:     actions.ActionArgString(string(sharedPoolStr)),
			workflow.ComponentIdsKey:   actions.ActionArgString(strings.Join(tidbGroupIds, ",")),
			workflow.ComponentTypesKey: actions.ActionArgString(strings.ToLower(apiv1.ComponentType_TIDB.String())),
			workflow.TenantIdKey:       actions.ActionArgString(tidbCluster.GetTenantId()),
			workflow.TiDBClusterIdKey:  actions.ActionArgString(tidbCluster.GetId()),
			utils.TraceIdKey:           utils.GetTraceId(ctx),
		}

		// set update options
		for k, v := range updateMap {
			args[k] = actions.ActionArgString(v)
		}

		dagExecId = utils.GeneratedDagExecId(ctx, fmt.Sprintf("set-tidbcluster-auditlog-%s-%s-%s",
			tidbCluster.GetSharedPoolId(),
			tidbCluster.GetId(),
			utils.FnvHash(string(sharedPoolStr)),
		))
		err = s.dagExecutor.Submit(ctx, dag.PlanUpdateTiDBCluster, args, dagExecId)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to submit create cluster dag: %s", err)
		}
		dagSubmitted = true

		return nil
	}); err != nil {
		if dagSubmitted {
			if abortErr := s.dagExecutor.Abort(ctx, dagExecId); abortErr != nil {
				tilog.Error(ctx, "failed to abort dag after transaction failure",
					zap.String("dagExecId", dagExecId),
					zap.Error(abortErr))
			}
		}
		return nil, err
	}

	return auditLogSetting, nil
}