package azure

import (
	"context"

	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v6"
	"k8s.io/utils/ptr"
)

const (
	ConnectionStatePending      = "Pending"
	ConnectionStateApproved     = "Approved"
	ConnectionStateRejected     = "Rejected"
	ConnectionStateDisconnected = "Disconnected"
	ConnectionStateEmpty        = ""
)

const (
	AcceptDescription = "Approved by TiDB Cloud."
	RejectDescription = "Rejected by TiDB Cloud."
	RemoveDescription = "Removed by TiDB Cloud."
)

type PrivateLinkServiceClient interface {
	ListPrivateLinkEndpointConnections(ctx context.Context, resourceGroupName string, serviceName string) ([]*armnetwork.PrivateEndpointConnection, error)
	AcceptPrivateLinkEndpointConnection(ctx context.Context, resourceGroupName string, serviceName string, conn *armnetwork.PrivateEndpointConnection) error
	RejectPrivateLinkEndpointConnection(ctx context.Context, resourceGroupName string, serviceName string, conn *armnetwork.PrivateEndpointConnection) error
	RemovePrivateLinkEndpointConnection(ctx context.Context, resourceGroupName string, serviceName string, conn *armnetwork.PrivateEndpointConnection) error
}

type privateLinkServiceClient struct {
	SubscriptionID string
	cli            *armnetwork.PrivateLinkServicesClient
}

func NewPrivateLinkServiceClient(subscriptionID string) (PrivateLinkServiceClient, error) {
	cred, err := azidentity.NewDefaultAzureCredential(nil)
	if err != nil {
		return nil, err
	}
	networkClientFactory, err := armnetwork.NewClientFactory(subscriptionID, cred, nil)
	if err != nil {
		return nil, err
	}
	cli := networkClientFactory.NewPrivateLinkServicesClient()
	return &privateLinkServiceClient{
		SubscriptionID: subscriptionID,
		cli:            cli,
	}, nil
}

func (c *privateLinkServiceClient) ListPrivateLinkEndpointConnections(ctx context.Context, resourceGroupName string, serviceName string) ([]*armnetwork.PrivateEndpointConnection, error) {
	pager := c.cli.NewListPrivateEndpointConnectionsPager(resourceGroupName, serviceName, nil)
	conns := make([]*armnetwork.PrivateEndpointConnection, 0)
	for pager.More() {
		page, err := pager.NextPage(ctx)
		if err != nil {
			return conns, err
		}
		conns = append(conns, page.Value...)
	}
	return conns, nil
}

func (c *privateLinkServiceClient) AcceptPrivateLinkEndpointConnection(ctx context.Context, resourceGroupName string, serviceName string, conn *armnetwork.PrivateEndpointConnection) error {
	connName := ptr.Deref(conn.Name, "")
	_, err := c.cli.UpdatePrivateEndpointConnection(ctx, resourceGroupName, serviceName, connName, armnetwork.PrivateEndpointConnection{
		Name: ptr.To(connName),
		Properties: &armnetwork.PrivateEndpointConnectionProperties{
			PrivateEndpoint: &armnetwork.PrivateEndpoint{
				ID: conn.Properties.PrivateEndpoint.ID,
			},
			PrivateLinkServiceConnectionState: &armnetwork.PrivateLinkServiceConnectionState{
				Description: ptr.To(AcceptDescription),
				Status:      ptr.To("Approved"),
			},
		},
	}, nil)
	if err != nil {
		return err
	}
	return nil
}

func (c *privateLinkServiceClient) RejectPrivateLinkEndpointConnection(ctx context.Context, resourceGroupName string, serviceName string, conn *armnetwork.PrivateEndpointConnection) error {
	connName := ptr.Deref(conn.Name, "")
	_, err := c.cli.UpdatePrivateEndpointConnection(ctx, resourceGroupName, serviceName, connName, armnetwork.PrivateEndpointConnection{
		Name: ptr.To(connName),
		Properties: &armnetwork.PrivateEndpointConnectionProperties{
			PrivateEndpoint: &armnetwork.PrivateEndpoint{
				ID: conn.Properties.PrivateEndpoint.ID,
			},
			PrivateLinkServiceConnectionState: &armnetwork.PrivateLinkServiceConnectionState{
				Description: ptr.To(RejectDescription),
				Status:      ptr.To("Rejected"),
			},
		},
	}, nil)
	if err != nil {
		return err
	}
	return nil
}

func (c *privateLinkServiceClient) RemovePrivateLinkEndpointConnection(ctx context.Context, resourceGroupName string, serviceName string, conn *armnetwork.PrivateEndpointConnection) error {
	connName := ptr.Deref(conn.Name, "")
	_, err := c.cli.BeginDeletePrivateEndpointConnection(ctx, resourceGroupName, serviceName, connName, nil)
	return err
}