package wireguard

import (
	"context"
	"net"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	v1vpnconfig "edge-infra.dev/pkg/sds/remoteaccess/k8s/apis/vpnconfigs/v1"

	"k8s.io/apimachinery/pkg/api/errors"

	"sigs.k8s.io/controller-runtime/pkg/client"

	"edge-infra.dev/pkg/edge/k8objectsutils"
	"edge-infra.dev/pkg/lib/crypto"
	"edge-infra.dev/pkg/sds/remoteaccess/constants"
	secrets "edge-infra.dev/pkg/sds/remoteaccess/wireguard/secret"
)

const (
	SecretPrefix     string = "keys-"
	SecretExpires    string = "expires"
	SecretPublicKey  string = "publicKey"
	SecretPrivateKey string = "privateKey"
	SecretIPAddress  string = "ipAddress"
)

type Instance struct {
	IPAddress  net.IP
	PublicKey  string
	PrivateKey string
}

func (i *Instance) GetIPAddress() net.IP {
	return i.IPAddress
}

func (i *Instance) GetPublicKey() string {
	return i.PublicKey
}

func (i *Instance) GetPrivateKey() string {
	return i.PrivateKey
}

// Create a new wireguard instance and K8s secret
func New(ctx context.Context, c client.Client, name, clusterEdgeID string, vpnConfig *v1vpnconfig.VPNConfig) (*Instance, error) {
	wgKey, err := crypto.GenerateWireguardKeyPair()
	if err != nil {
		return nil, err
	}
	instance := &Instance{
		PublicKey:  wgKey.PublicKey(),
		PrivateKey: wgKey.PrivateKey(),
	}
	if err := createInstanceSecret(ctx, c, name, clusterEdgeID, instance, vpnConfig); err != nil {
		return nil, err
	}
	return instance, nil
}

// Upon updating the IP address we update the K8s secret
func (i *Instance) UpdateIPAddress(ctx context.Context, c client.Client, name, clusterEdgeID string, ip net.IP) error {
	i.IPAddress = ip
	return updateInstanceSecret(ctx, c, name, clusterEdgeID, i)
}

// Checks for existing wireguard instance secret
// If no secret exists or has expired, create new instance, marshal to json, save as K8s secret, return instance
// If secret exists, unmarshal secret and return instance
func GetInstance(ctx context.Context, c client.Client, name, clusterEdgeID string, vpnConfig *v1vpnconfig.VPNConfig) (*Instance, error) {
	secret := &corev1.Secret{}
	secretName := k8objectsutils.NameWithPrefix(SecretPrefix+name, clusterEdgeID)

	err := c.Get(ctx, client.ObjectKey{Namespace: constants.VPNNamespace, Name: secretName}, secret)
	if err != nil && !errors.IsNotFound(err) {
		return nil, err
	}
	if errors.IsNotFound(err) {
		return New(ctx, c, name, clusterEdgeID, vpnConfig)
	}

	// check if secret has expired
	expires, err := time.Parse(time.RFC3339, secret.Annotations[SecretExpires])
	if err != nil {
		return nil, err
	}
	if time.Now().After(expires) {
		return rotateInstanceSecret(ctx, c, name, clusterEdgeID, secret, vpnConfig)
	}

	return instanceFromSecret(secret), nil
}

func rotateInstanceSecret(ctx context.Context, c client.Client, name, clusterEdgeID string, secret *corev1.Secret, vpnConfig *v1vpnconfig.VPNConfig) (*Instance, error) {
	// we need to ensure that we use the previously set IP address when rotating instance secrets
	old := instanceFromSecret(secret)

	if err := c.Delete(ctx, secret); err != nil {
		return nil, err
	}

	instance, err := New(ctx, c, name, clusterEdgeID, vpnConfig)
	if err != nil {
		return nil, err
	}

	err = instance.UpdateIPAddress(ctx, c, name, clusterEdgeID, old.IPAddress)
	return instance, err
}

func instanceFromSecret(secret *corev1.Secret) *Instance {
	return &Instance{
		PublicKey:  string(secret.Data[SecretPublicKey]),
		PrivateKey: string(secret.Data[SecretPrivateKey]),
		IPAddress:  net.ParseIP(string(secret.Data[SecretIPAddress])),
	}
}

func createInstanceSecret(ctx context.Context, c client.Client, name, clusterEdgeID string, instance *Instance, vpnConfig *v1vpnconfig.VPNConfig) error {
	secretName := k8objectsutils.NameWithPrefix(SecretPrefix+name, clusterEdgeID)
	secret := &corev1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      secretName,
			Namespace: constants.VPNNamespace,
			Annotations: map[string]string{
				SecretExpires: secrets.ExpireAt().Format(time.RFC3339),
			},
			OwnerReferences: createOwnerReferences(name, vpnConfig),
		},
		Data: map[string][]byte{
			SecretPublicKey:  []byte(instance.GetPublicKey()),
			SecretPrivateKey: []byte(instance.GetPrivateKey()),
			SecretIPAddress:  []byte(instance.GetIPAddress().String()),
		},
	}
	return c.Create(ctx, secret)
}

func updateInstanceSecret(ctx context.Context, c client.Client, name, clusterEdgeID string, instance *Instance) error {
	secretName := k8objectsutils.NameWithPrefix(SecretPrefix+name, clusterEdgeID)
	secret := &corev1.Secret{}
	if err := c.Get(ctx, client.ObjectKey{Namespace: constants.VPNNamespace, Name: secretName}, secret); err != nil {
		return err
	}
	secret.Data[SecretPublicKey] = []byte(instance.GetPublicKey())
	secret.Data[SecretPrivateKey] = []byte(instance.GetPrivateKey())
	secret.Data[SecretIPAddress] = []byte(instance.GetIPAddress().String())
	return c.Update(ctx, secret)
}

func createOwnerReferences(name string, vpnConfig *v1vpnconfig.VPNConfig) []metav1.OwnerReference {
	if name != constants.StoreName || vpnConfig == nil {
		return []metav1.OwnerReference{}
	}
	return []metav1.OwnerReference{
		{
			APIVersion: vpnConfig.APIVersion,
			Kind:       vpnConfig.Kind,
			Name:       vpnConfig.GetName(),
			UID:        vpnConfig.GetUID(),
		},
	}
}