package trustanchor

import (
	"context"
	"crypto/x509"
	"fmt"
	"slices"
	"time"

	"github.com/linkerd/linkerd2/pkg/tls"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"edge-infra.dev/pkg/edge/linkerd"
	l5dv1alpha1 "edge-infra.dev/pkg/edge/linkerd/k8s/apis/linkerd/v1alpha1"
	"edge-infra.dev/pkg/edge/linkerd/k8s/controllers/metrics"
	"edge-infra.dev/pkg/k8s/runtime/conditions"
	"edge-infra.dev/pkg/lib/fog"
)

const (
	CaBundleKey               = "ca-bundle.crt"
	ManualTrustAnchorRotation = "linkerd.io/manual-anchor-rotation"
	trustAnchorRotated        = "linkerd.io/trust-anchor-rotated"
)

// GenerateTrustAnchor creates the trust anchor cert and key
// https://linkerd.io/2.11/tasks/automatically-rotating-control-plane-tls-credentials/#save-the-signing-key-pair-as-a-secret
func GenerateTrustAnchor(ctx context.Context) ([]byte, []byte, error) {
	log := fog.FromContext(ctx).WithName("generate-trust-anchor")

	// generate private key for root ca
	rootCAkey, err := tls.GenerateKey()
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate key for root ca: %w", err)
	}

	// creating root certificate authority with a duration of two years
	certExpiryDate := time.Now().AddDate(linkerd.CertDurationYear, 0, 0)
	metrics.RecordTrustAnchorExpiryTime(float64(certExpiryDate.Unix()))

	tlsValidity := tls.Validity{
		Lifetime: time.Until(certExpiryDate),
	}
	ca, err := tls.CreateRootCA("root.linkerd.cluster.local", rootCAkey, tlsValidity)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate root ca: %w", err)
	}

	cert := ca.Cred.Crt.EncodePEM()
	key, err := tls.EncodePrivateKeyPEM(rootCAkey)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to encode private key: %w", err)
	}
	log.Info("generated a trust anchor certificate and key", "certificate expiry date", certExpiryDate)
	return []byte(cert), key, nil
}

func buildSecret(l5d *l5dv1alpha1.Linkerd, cert, key []byte) *corev1.Secret {
	return &corev1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:            linkerd.TrustAnchorName,
			Namespace:       linkerd.Namespace,
			OwnerReferences: linkerd.OwnerRef(l5d),
		},
		Type: corev1.SecretTypeTLS,
		Data: map[string][]byte{
			corev1.TLSCertKey:       cert,
			corev1.TLSPrivateKeyKey: key,
		},
	}
}

// CreateIfNotExists checks if the trust anchor secret has already been
// created and will create it if it isn't present.
//
// The content of the trust anchor secret that is required by the linkerd
// installation manifest rendering process is returned regardless so that manifests
// can be rendered correctly after controller restarts without re-generating the
// secret each time.
func CreateIfNotExists(ctx context.Context, c client.Client, l5d *l5dv1alpha1.Linkerd) (string, error) {
	cert, err := SecretExists(ctx, c)
	if err != nil {
		conditions.MarkFalse(l5d, l5dv1alpha1.TrustAnchor, l5dv1alpha1.TrustAnchorSecretSetupFailedReason, "%v", err)
		return "", err
	}

	// if cert exists then return early
	if cert != "" {
		return cert, nil
	}

	return Create(ctx, c, l5d)
}

// Create is responsible for creating a new trust anchor secret
func Create(ctx context.Context, c client.Client, l5d *l5dv1alpha1.Linkerd) (string, error) {
	cert, err := createSecret(ctx, c, l5d)
	if err != nil {
		conditions.MarkFalse(l5d, l5dv1alpha1.TrustAnchor, l5dv1alpha1.TrustAnchorSecretSetupFailedReason, "%v", err)
		return "", err
	}

	conditions.MarkTrue(l5d, l5dv1alpha1.TrustAnchor, l5dv1alpha1.SucceededReason, "Created trust anchor secret")
	return cert, nil
}

// Rotate generates a new trust anchor secret, annotates it, updates the ca bundle and patches the existing secret
func Rotate(ctx context.Context, c client.Client, l5d *l5dv1alpha1.Linkerd) error {
	log := fog.FromContext(ctx).WithName("rotate")

	newTrustAnchorSecret, newSecret, err := generateSecret(ctx, l5d)
	if err != nil {
		log.Error(err, "failed to generate trust anchor secret")
		return err
	}
	if newSecret.Annotations == nil {
		newSecret.Annotations = make(map[string]string)
	}
	// annotate secret to prevent rotation on next reconcile
	newSecret.Annotations[trustAnchorRotated] = "true"

	oldSecret, err := getSecret(ctx, c)
	if err != nil {
		return err
	} else if oldSecret == nil {
		log.Info("trust anchor secret doesn't exist - creating secret")
		if err := c.Create(ctx, newSecret, linkerd.CreateOpts()); err != nil {
			return err
		}
	}
	oldTrustAnchorCert, err := getCert(oldSecret)
	if err != nil {
		log.Error(err, "failed to get old trust anchor certificate from secret")
		return err
	}

	if _, err := updateCaBundle(ctx, c, oldTrustAnchorCert, newTrustAnchorSecret); err != nil {
		log.Error(err, "failed updating the ca bundle in the identity configmap")
		return err
	}
	log.Info("successfully updated the identity configmap with the new ca bundle")

	secret, err := getSecret(ctx, c)
	if err != nil {
		return err
	} else if secret == nil {
		log.Info("trust anchor secret doesn't exist - creating secret")
		conditions.MarkTrue(l5d, l5dv1alpha1.TrustAnchor, l5dv1alpha1.SucceededReason, "Created trust anchor secret")
		return c.Create(ctx, newSecret, linkerd.CreateOpts())
	}

	// if trust anchor secret did not exist prior to rotation, one will have been created
	// so no need to patch
	if oldSecret != nil {
		conditions.MarkTrue(l5d, l5dv1alpha1.TrustAnchor, l5dv1alpha1.DualAnchor, "successfully merged old and new secret")
		return c.Patch(ctx, newSecret, client.StrategicMergeFrom(oldSecret.DeepCopy()))
	}

	conditions.MarkTrue(l5d, l5dv1alpha1.TrustAnchor, l5dv1alpha1.SucceededReason, "Created trust anchor secret")
	return nil
}

// IsRotated checks for the trust anchor rotated annotation on the trust
// anchor secret and returns true if it is.
func IsRotated(ctx context.Context, c client.Client) bool {
	log := fog.FromContext(ctx).WithName("is-rotated")

	secret, err := getSecret(ctx, c)
	if err != nil {
		log.Error(err, "failed to get secret")
		return false
	}
	if secret == nil {
		return false
	}

	_, annotationExists := secret.Annotations[trustAnchorRotated]
	if !annotationExists {
		log.Info("trust anchor rotation annotation does not exist - trust anchor has not been rotated")
		return false
	}

	log.Info("the trust anchor has been rotated")
	return true
}

func HasManualRotationAnnotation(l5d *l5dv1alpha1.Linkerd) bool {
	_, exists := l5d.Annotations[ManualTrustAnchorRotation]
	return exists
}

func RemoveRotationAnnotations(ctx context.Context, c client.Client, l5d *l5dv1alpha1.Linkerd) error {
	secret, err := getSecret(ctx, c)
	if err != nil {
		return err
	}
	secretPatch := []byte(fmt.Sprintf(
		`{"metadata":{"annotations":{"%s":null}}}`,
		trustAnchorRotated),
	)
	if err := c.Patch(ctx, secret, client.RawPatch(types.MergePatchType, secretPatch)); err != nil {
		return err
	}

	l5dPatch := []byte(fmt.Sprintf(
		`{"metadata":{"annotations":{"%s":null}}}`,
		ManualTrustAnchorRotation),
	)
	return c.Patch(ctx, l5d.DeepCopy(), client.RawPatch(types.MergePatchType, l5dPatch))
}

// SecretExists attempts to get the trust anchor secret and return the cert string if it exists.
// Returning an empty string and nil error indicates the secret does not exist.
func SecretExists(ctx context.Context, c client.Client) (string, error) {
	secret, err := getSecret(ctx, c)
	if err != nil {
		return "", err
	} else if secret == nil {
		return "", nil
	}

	return getCert(secret)
}

func getCert(secret *corev1.Secret) (string, error) {
	if secret == nil {
		return "", nil
	}
	data, ok := secret.Data[corev1.TLSCertKey]
	if !ok || len(data) == 0 {
		return "", fmt.Errorf("secret was present, but %s was not present or empty", corev1.TLSCertKey)
	}
	return string(secret.Data[corev1.TLSCertKey]), nil
}

// UpdateCaBundle checks if the ca bundle has the current trust anchor
// certificate and updates the bundle if not.
func UpdateCaBundle(ctx context.Context, c client.Client) (string, error) {
	log := fog.FromContext(ctx).WithName("update-ca-bundle")
	caBundle, err := GetCaBundle(ctx, c)
	if err != nil {
		return "", err
	}

	if caBundleOK, err := CheckCaBundle(ctx, c); err != nil || !caBundleOK {
		if err != nil {
			return "", err
		}
		log.Info("ca bundle is out-of-date - updating...")

		lastCa, err := getLastValidCa(caBundle)
		if err != nil {
			return "", err
		}

		newCa, err := SecretExists(ctx, c)
		if err != nil {
			return "", err
		}
		caBundle, err = updateCaBundle(ctx, c, lastCa, newCa)
		if err != nil {
			return "", err
		}
	}

	return caBundle, nil
}

// updateCaBundle patches the ca bundle in the identity configmap with the new trust anchor secret
func updateCaBundle(ctx context.Context, c client.Client, oldTrustAnchorSecret, newTrustAnchorSecret string) (string, error) {
	updatedCaBundle := fmt.Sprintf("%s%s", oldTrustAnchorSecret, newTrustAnchorSecret)

	cm, err := getIdentityCM(ctx, c)
	if err != nil {
		return updatedCaBundle, err
	}
	if cm == nil {
		return updatedCaBundle, fmt.Errorf("unable to update the ca bundle in the identity configmap")
	}

	cmPatch := cm.DeepCopy()
	if cmPatch.Data == nil {
		cmPatch.Data = make(map[string]string)
	}
	cmPatch.Data[CaBundleKey] = updatedCaBundle

	return updatedCaBundle, c.Patch(ctx, cmPatch, client.StrategicMergeFrom(cm.DeepCopy()))
}

// GetCaBundle retrieves the ca bundle from the linkerd-identity-trust-roots configmap
// or the trust anchor secret if the configmap does not exist
func GetCaBundle(ctx context.Context, c client.Client) (string, error) {
	log := fog.FromContext(ctx).WithName("get-ca-bundle")
	const getCaFailMsg string = "unable to get the ca bundle"

	cm, err := getIdentityCM(ctx, c)
	if errors.IsNotFound(err) {
		log.Info("identity configmap not found - get ca from the trust anchor secret")
		caBundle, err := SecretExists(ctx, c)

		if err == nil && caBundle == "" {
			missingTrustErr := fmt.Errorf("trustanchor secret does not exist")
			log.Error(missingTrustErr, getCaFailMsg)
			return "", missingTrustErr
		} else if err != nil {
			log.Error(err, getCaFailMsg)
			return "", err
		}

		return caBundle, nil
	} else if err != nil {
		log.Error(err, getCaFailMsg)
		return "", err
	}

	bundle, exists := cm.Data[CaBundleKey]
	if !exists {
		missingCABundleErr := fmt.Errorf("ca bundle not found in configmap")
		log.Error(missingCABundleErr, getCaFailMsg)
		return "", missingCABundleErr
	}
	return bundle, nil
}

// CheckCaBundle checks that the ca bundle in the identity configmap contains the currently
// deployed trust anchor secret.
func CheckCaBundle(ctx context.Context, c client.Client) (bool, error) {
	log := fog.FromContext(ctx).WithName("check-ca-bundle")
	currentRootCa, err := SecretExists(ctx, c)
	if err != nil {
		return false, err
	} else if currentRootCa == "" {
		return false, fmt.Errorf("trust anchor secret does not exist")
	}

	caBundle, err := GetCaBundle(ctx, c)
	if err != nil {
		return false, err
	}

	bundleCerts, err := tls.DecodePEMCertificates(caBundle)
	if err != nil {
		return false, err
	}
	currentCert, err := tls.DecodePEMCertificates(currentRootCa)
	if err != nil {
		return false, err
	}
	if slices.ContainsFunc(currentCert, func(entry *x509.Certificate) bool {
		for _, cert := range bundleCerts {
			if cert.Equal(entry) {
				return true
			}
		}
		return false
	}) {
		return true, nil
	}

	log.Info("identity configmap is not up-to-date")
	return false, nil
}

// getLastValidCa returns the ca in the bundle that was created last
func getLastValidCa(caBundle string) (string, error) {
	certs, err := tls.DecodePEMCertificates(caBundle)
	if err != nil {
		return "", err
	}

	if len(certs) == 0 {
		return "", nil
	} else if len(certs) > 1 {
		// sort ca bundle certificates in ascending order from the earliest
		// validity to the latest validity
		slices.SortFunc(certs, func(certA, certB *x509.Certificate) int {
			if certA.NotAfter.Before(certB.NotAfter) {
				return -1
			}
			return 1
		})
	}

	return tls.EncodeCertificatesPEM(certs[len(certs)-1]), nil
}

func getIdentityCM(ctx context.Context, c client.Client) (*corev1.ConfigMap, error) {
	cm := &corev1.ConfigMap{}
	err := c.Get(ctx, types.NamespacedName{Name: linkerd.LinkerdIdentityConfigMap, Namespace: linkerd.Namespace}, cm)
	if err != nil {
		return nil, err
	}
	return cm, nil
}

func getSecret(ctx context.Context, c client.Client) (*corev1.Secret, error) {
	log := fog.FromContext(ctx).WithName("get-secret")
	secret := &corev1.Secret{}
	err := c.Get(ctx, linkerd.TrustAnchorKey(), secret)
	if errors.IsNotFound(err) {
		log.Info("the trust anchor secret was not found")
		return nil, nil
	} else if err != nil {
		log.Error(err, "failed to get trust anchor secret")
		return nil, fmt.Errorf("failed to check if trust anchor secret exists: %w", err)
	}
	return secret, err
}

// createSecret generates the Linkerd trust anchor signing key pair
// and creates the K8s secret Linkerd is expecting if it does not already exist.
// https://linkerd.io/2.11/tasks/automatically-rotating-control-plane-tls-credentials/#save-the-signing-key-pair-as-a-secret
func createSecret(ctx context.Context, c client.Client, l5d *l5dv1alpha1.Linkerd) (string, error) {
	log := fog.FromContext(ctx).WithName("create-secret")
	cert, secret, err := generateSecret(ctx, l5d)
	if err != nil {
		log.Error(err, "failed to generate trust anchor secret")
		return "", err
	}

	if err := c.Create(ctx, secret, linkerd.CreateOpts()); client.IgnoreAlreadyExists(err) != nil {
		log.Error(err, "failed to create trust anchor secret")
		return "", fmt.Errorf("failed to create secret")
	}
	log.Info("the trust anchor secret was created or already exists")
	return cert, nil
}

func generateSecret(ctx context.Context, l5d *l5dv1alpha1.Linkerd) (string, *corev1.Secret, error) {
	cert, key, err := GenerateTrustAnchor(ctx)
	if err != nil {
		return "", nil, err
	}
	return string(cert), buildSecret(l5d, cert, key), nil
}