package action

import (
	"context"
	"fmt"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	"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"
)

// AWS NLB tuning annotation keys passed through DPM Spec.Annotations. Caller
// values win over DPM controller defaults (see
// internal/dataplane-manager/pkg/controllers/loadbalancer/task/eks/service.go
// `overlayPrivateSvc`), which is exactly what we want for the cassandra-proxy
// healthcheck endpoint and the custom target-group-attributes.
const (
	awsLBHealthcheckPathKey       = "service.beta.kubernetes.io/aws-load-balancer-healthcheck-path"
	awsLBHealthcheckPortKey       = "service.beta.kubernetes.io/aws-load-balancer-healthcheck-port"
	awsLBHealthcheckProtocolKey   = "service.beta.kubernetes.io/aws-load-balancer-healthcheck-protocol"
	awsLBTargetGroupAttributesKey = "service.beta.kubernetes.io/aws-load-balancer-target-group-attributes"
)

// Cassandra-proxy NLB listener / healthcheck constants. These mirror
// `tasks/cassandra-proxy-on-premium/dev/all-in-one.yaml`. Held as constants
// rather than tunables because changing them would invalidate existing
// customer endpoints.
const (
	cassandraProxyCQLPort               = int32(9042)
	cassandraProxyCQLSSLPort            = int32(9142)
	cassandraProxyHealthcheckPath       = "/ready"
	cassandraProxyHealthcheckPort       = "6060"
	cassandraProxyHealthcheckProtocol   = "http"
	cassandraProxyTargetGroupAttributes = "preserve_client_ip.enabled=true,stickiness.enabled=false"
	cassandraProxyCQLPortName           = "cassandra"
	cassandraProxyCQLSSLPortName        = "cassandra-ssl"
	cassandraProxyListenerProtocol      = "TCP"
)

const ApplyCassandraProxyLBActionName = "apply-cassandra-proxy-lb"

type ApplyCassandraProxyLBAction struct {
	sharedPoolRepo  repo.SharedPoolRepo
	tidbClusterRepo repo.TiDBClusterRepo
	lbClient        dplanemgrv1.LoadBalancerServiceClient
}

func NewApplyCassandraProxyLBAction(
	sharedPoolRepo repo.SharedPoolRepo,
	tidbClusterRepo repo.TiDBClusterRepo,
	lbClient dplanemgrv1.LoadBalancerServiceClient,
) actions.Action {
	return &ApplyCassandraProxyLBAction{
		sharedPoolRepo:  sharedPoolRepo,
		tidbClusterRepo: tidbClusterRepo,
		lbClient:        lbClient,
	}
}

func (a *ApplyCassandraProxyLBAction) Name() string {
	return ApplyCassandraProxyLBActionName
}

// Invoke creates or updates the cassandra-proxy LoadBalancer CR for every
// CASSANDRA_PROXY component named in `ComponentIdsKey`. This intentionally
// mirrors ApplyLoadBalanceAction: Get -> Create if NotFound, otherwise Update.
func (a *ApplyCassandraProxyLBAction) Invoke(args actions.ActionArgs) error {
	ctx := utils.GenerateContext(args)
	tilog.Info(ctx, "apply cassandra-proxy lb action")

	if _, err := args.GetString(workflow.ComponentIdsKey); err != nil {
		return fmt.Errorf("target component id is required, but not found")
	}

	sharedPoolId, err := args.GetString(workflow.SharedPoolIdKey)
	if err != nil {
		return fmt.Errorf("failed to get sharedPool id: %w", err)
	}

	sharedPool, err := a.sharedPoolRepo.Load(ctx, sharedPoolId, false)
	if err != nil {
		return fmt.Errorf("failed to load sharedPool: %w", err)
	}

	// Frozen-snapshot semantics: read ComponentGroups from the args payload
	// captured at DAG-submit time. Same shape as apply_cassandra_proxy_group.go.
	sharedPoolStr, err := args.GetString(workflow.SharedPoolKey)
	if err != nil {
		return fmt.Errorf("failed to get sharedPool: %w", err)
	}
	var tempSharedPool entity.SharedPool
	if err := protojson.Unmarshal([]byte(sharedPoolStr), &tempSharedPool); err != nil {
		return fmt.Errorf("failed to unmarshal sharedPool: %w", err)
	}
	sharedPool.ComponentGroups = tempSharedPool.GetComponentGroups()

	for _, comp := range sharedPool.GetComponentGroups() {
		if comp.GetComponentType() != apiv1.ComponentType_CASSANDRA_PROXY {
			continue
		}
		if !workflow.CheckIsTargetComponent(ctx, comp, args) {
			continue
		}

		tidbCluster, err := a.tidbClusterRepo.Load(ctx, comp.GetId(), false)
		if err != nil {
			return fmt.Errorf("failed to load tidb cluster %s: %w", comp.GetId(), err)
		}
		if tidbCluster.GetGatewayIdentity() == "" {
			return status.Errorf(codes.FailedPrecondition,
				"tidb cluster %s has empty gateway_identity; required for LB Identity",
				comp.GetId())
		}

		desired := buildCassandraProxyLB(comp, tidbCluster, sharedPool)
		if err := a.reconcileLB(ctx, sharedPool.GetDataplaneName(), desired); err != nil {
			return err
		}
	}

	return nil
}

// reconcileLB follows the TiProxy LB pattern: create missing CRs and update
// existing CRs with the desired full spec.
func (a *ApplyCassandraProxyLBAction) reconcileLB(
	ctx context.Context,
	dataplaneId string,
	desired *dplanemgrv1.LoadBalancer,
) error {
	_, err := a.lbClient.GetLoadBalancer(ctx, &dplanemgrv1.GetLoadBalancerRequest{
		DataPlaneId: dataplaneId,
		Id:          desired.GetId(),
	})
	switch {
	case err == nil:
		return a.updateLB(ctx, dataplaneId, desired)

	case status.Code(err) == codes.NotFound:
		_, createErr := a.lbClient.CreateLoadBalancer(ctx, &dplanemgrv1.CreateLoadBalancerRequest{
			DataPlaneId:  dataplaneId,
			Id:           desired.GetId(),
			Loadbalancer: desired,
		})
		if createErr == nil {
			return nil
		}
		// TOCTOU: another reconciler created it between our Get and Create.
		// Fall back to Update so we still settle on `desired`.
		if status.Code(createErr) == codes.AlreadyExists {
			tilog.Info(ctx, "cassandra-proxy LB exists after Create AlreadyExists, updating",
				zap.String("lbName", desired.GetId()))
			return a.updateLB(ctx, dataplaneId, desired)
		}
		return fmt.Errorf("failed to create cassandra-proxy LB %s: %w", desired.GetId(), createErr)

	default:
		return fmt.Errorf("failed to get cassandra-proxy LB %s: %w", desired.GetId(), err)
	}
}

func (a *ApplyCassandraProxyLBAction) updateLB(
	ctx context.Context,
	dataplaneId string,
	desired *dplanemgrv1.LoadBalancer,
) error {
	if _, err := a.lbClient.UpdateLoadBalancer(ctx, &dplanemgrv1.UpdateLoadBalancerRequest{
		DataPlaneId:  dataplaneId,
		Id:           desired.GetId(),
		Loadbalancer: desired,
	}); err != nil {
		return fmt.Errorf("failed to update cassandra-proxy LB %s: %w", desired.GetId(), err)
	}
	return nil
}

// buildCassandraProxyLB returns the desired LoadBalancer CR. The fields
// listed here are the contract this action owns; output-only fields (DNS,
// Phase, PrivateIps) are populated by DPM.
func buildCassandraProxyLB(
	comp *entity.Component,
	tidbCluster *entity.TiDBCluster,
	sharedPool *entity.SharedPool,
) *dplanemgrv1.LoadBalancer {
	return &dplanemgrv1.LoadBalancer{
		Id:          entity.GetCassandraProxyLbName(comp.GetId()),
		Identity:    tidbCluster.GetGatewayIdentity(),
		SubIdentity: entity.BuildCassandraProxySubIdentity(),
		Namespace:   entity.GetDataPlaneNamespace(sharedPool.GetId()),
		Type:        dplanemgrv1.LoadBalancer_PRIVATE,
		// Legacy single-port field. Newer DPM controllers prefer Ports, but
		// still tolerate Port; keep it filled as defense in depth because the
		// capability probe already rejects dataplanes that do not support Ports.
		Port: cassandraProxyCQLPort,
		Ports: []*dplanemgrv1.LoadBalancer_Port{
			{Name: cassandraProxyCQLPortName, Port: cassandraProxyCQLPort, TargetPort: cassandraProxyCQLPort, Protocol: cassandraProxyListenerProtocol},
			{Name: cassandraProxyCQLSSLPortName, Port: cassandraProxyCQLSSLPort, TargetPort: cassandraProxyCQLSSLPort, Protocol: cassandraProxyListenerProtocol},
		},
		Annotations: map[string]string{
			awsLBHealthcheckPathKey:       cassandraProxyHealthcheckPath,
			awsLBHealthcheckPortKey:       cassandraProxyHealthcheckPort,
			awsLBHealthcheckProtocolKey:   cassandraProxyHealthcheckProtocol,
			awsLBTargetGroupAttributesKey: cassandraProxyTargetGroupAttributes,
		},
		Selectors: entity.BuildCassandraProxySelector(comp.GetId()),
		Tags:      entity.BuildCassandraProxyTags(comp, tidbCluster, sharedPool),
	}
}