package bannerctl

import (
	"bytes"
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"database/sql"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"math/big"
	"path/filepath"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	bannerAPI "edge-infra.dev/pkg/edge/apis/banner/v1alpha1"
	sequelApi "edge-infra.dev/pkg/edge/apis/sequel/k8s/v1alpha2"
	"edge-infra.dev/pkg/edge/constants"
	workloadApi "edge-infra.dev/pkg/edge/constants/api/workload"
	"edge-infra.dev/pkg/k8s/runtime/controller/reconcile/recerr"
	secretMgrApi "edge-infra.dev/pkg/lib/gcp/secretmanager"
	"edge-infra.dev/pkg/lib/uuid"
)

const (
	createCAPool   = "INSERT INTO ca_pools (banner_edge_id) VALUES ($1) RETURNING ca_pool_edge_id"
	createCACert   = "INSERT INTO ca_certificates (ca_pool_edge_id, status, cert_ref, private_key_ref, expiration) VALUES ($1, $2, $3, $4, $5) RETURNING ca_cert_edge_id"
	getActiveCert  = "SELECT ca_cert_edge_id, expiration FROM ca_certificates WHERE ca_pool_edge_id = $1 AND status = 'active'"
	getStagedCert  = "SELECT ca_cert_edge_id, expiration FROM ca_certificates WHERE ca_pool_edge_id = $1 AND status = 'staged'"
	getRetiredCert = "SELECT ca_cert_edge_id, expiration FROM ca_certificates WHERE ca_pool_edge_id = $1 AND status = 'retired'"
	getCAPoolID    = `SELECT ca_pool_edge_id FROM ca_pools WHERE banner_edge_id = $1`
	staged         = "staged"
	active         = "active"
	retired        = "retired"
	deleted        = "deleted"
)

// reconcileCerts reconciles the CA pool and CA certs for a banner.
func (r *BannerReconciler) reconcileCerts(ctx context.Context, banner *bannerAPI.Banner) recerr.Error {
	log := r.Log.WithValues("banner", banner.Name)

	// check or create CA pool per banner
	poolID, err := r.reconcileCAPool(ctx, banner.Name)
	if err != nil {
		return recerr.New(err, bannerAPI.PlatformSecretsCreationFailedReason)
	}

	// check or create CA cert per banner
	err = r.reconcileCACerts(ctx, banner, poolID)
	if err != nil {
		return recerr.New(err, bannerAPI.PlatformSecretsCreationFailedReason)
	}

	log.Info("Successfully reconciled CA pools and certs for banner", "banner", banner.Name)
	return nil
}

// reconcileCAPool reconciles the CA pool for a banner.
func (r *BannerReconciler) reconcileCAPool(ctx context.Context, bannerID string) (string, error) {
	tx, err := r.EdgeDB.BeginTx(ctx, &sql.TxOptions{})
	if err != nil {
		return "", err
	}

	defer func() {
		if err != nil {
			err = errors.Join(err, tx.Rollback())
		}
	}()

	log := r.Log.WithValues("banner", bannerID)

	var poolID string
	err = tx.QueryRowContext(ctx, getCAPoolID, bannerID).Scan(&poolID)
	if err != nil && err != sql.ErrNoRows {
		return "", fmt.Errorf("failed to retrieve CA pool: %w", err)
	} else if err == nil && poolID != "" {
		log.Info("CA pool already exists", "caPool", poolID)
		return poolID, nil
	}

	err = tx.QueryRowContext(ctx, createCAPool, bannerID).Scan(&poolID)
	if err != nil {
		return "", fmt.Errorf("failed to insert CA pool into database: %w", err)
	}

	if err = tx.Commit(); err != nil {
		return "", err
	}

	log.Info("Created CA pool", "caPool", poolID)
	return poolID, nil
}

// reconcileCACerts reconciles the CA certs for a banner.
// When a CA is approaching expiry (1 year left), a new CA cert is generated, with status of staged
// After 3 months (tbd) of the new CA cert existing, the old CA cert is set to retired and the new CA cert is marked as active
// If the CA cert is expired, it is marked as deleted
func (r *BannerReconciler) reconcileCACerts(ctx context.Context, b *bannerAPI.Banner, caPool string) error {
	tx, err := r.EdgeDB.BeginTx(ctx, &sql.TxOptions{})
	if err != nil {
		return err
	}

	defer func() {
		if err != nil {
			err = errors.Join(err, tx.Rollback())
		}
	}()

	log := r.Log.WithValues("banner", b.Name)

	// STEP 1 : Check if there is an active cert
	// if there is no active cert, check if there is a staged cert
	// if there is a staged cert make that active
	// if there is no staged cert, create a new active cert

	var activeCertID string
	var activeCertExpiration time.Time
	var stagedCertID string
	var stagedCertExpiration time.Time
	// This will retrieve the active cert for a given CA pool edge ID
	// if there is no active cert, it will create one
	err = tx.QueryRowContext(ctx, getActiveCert, caPool).Scan(&activeCertID, &activeCertExpiration)
	if err != nil && err != sql.ErrNoRows {
		return fmt.Errorf("failed to query active CA cert: %w", err)
	} else if err == sql.ErrNoRows {
		// if there is no active but there is a staged then we want to make that active
		err = tx.QueryRowContext(ctx, getStagedCert, caPool).Scan(&stagedCertID, &activeCertExpiration)
		if err != nil && err != sql.ErrNoRows {
			return fmt.Errorf("failed to query staged CA cert: %w", err)
		} else if err == sql.ErrNoRows {
			// if there is no active and no staged then we want to create a new active
			return r.createCA(ctx, b, caPool, active)
		}
		// if there is no errors then we want to make the staged cert active
		// this should ideally not be used, but is a catch for if the one cycle rotation fails for any reason
		err = updateCertificateStatus(ctx, tx, stagedCertID, active)
		if err != nil {
			return fmt.Errorf("failed to update staged certificate status: %w", err)
		}
	}

	//STEP 2 : Check if there is a staged cert

	var stagedNeeded bool
	var stagedExists bool
	// This will retrieve the staged cert for a given CA pool edge ID
	// if there is a staged cert make stagedExists true so we dont create another one
	err = tx.QueryRowContext(ctx, getStagedCert, caPool).Scan(&stagedCertID, &stagedCertExpiration)
	if err != nil && err != sql.ErrNoRows {
		return fmt.Errorf("failed to query staged CA cert: %w", err)
	} else if err == sql.ErrNoRows {
		stagedExists = false
	} else {
		stagedExists = true
	}

	// STEP 3 : Check if the active cert is expired or close to expiration
	// if it is close to expiration, check if there is a staged cert
	// if there is a staged cert, make the active cert retired and the staged cert active
	// if there is no staged cert, create a new staged cert

	// if the active cert is less than 9 months away from expiration and there is a staged cert
	// we will mark the active cert as retired and the staged cert as active
	// if the active cert is less than 1 year & 3 months away from expiration and there is no staged cert
	// we will mark stagedNeeded as true so that we can create a new staged cert
	if activeCertExpiration.Before(time.Now().AddDate(1, 3, 0)) &&
		activeCertExpiration.Before(time.Now().AddDate(0, 9, 0)) && stagedExists {
		//update the current active cert to retired
		err = updateCertificateStatus(ctx, tx, activeCertID, retired)
		if err != nil {
			return fmt.Errorf("failed to update certificate status: %w", err)
		}
		log.Info("Updated active cert to retired", "certID", activeCertID)
		//update the current staged cert to active
		//this allows us to perform the rotation in one cycle
		//if the above step fails for some reason then staged will become active on the next cycle
		err = updateCertificateStatus(ctx, tx, stagedCertID, active)
		if err != nil {
			return fmt.Errorf("failed to update certificate status: %w", err)
		}
		log.Info("Updated staged cert to active", "certID", stagedCertID)
	} else if activeCertExpiration.Before(time.Now().AddDate(1, 3, 0)) {
		stagedNeeded = true
	}

	// if we need a staged cert and one does not exist, create one
	if stagedNeeded && !stagedExists {
		err = r.createCA(ctx, b, caPool, staged)
		if err != nil {
			return fmt.Errorf("failed to create CA cert: %w", err)
		}
	}

	// STEP 4 : Check if the retired cert is expired
	// if it is expired, mark it as deleted

	var retiredCertID string
	var retiredCertExpiration time.Time
	// This will retrieve the retired cert for a given CA pool edge ID
	// if there is a retired cert and it is expired, it will be marked as deleted
	err = tx.QueryRowContext(ctx, getRetiredCert, caPool).Scan(&retiredCertID, &retiredCertExpiration)
	if err != nil && err != sql.ErrNoRows {
		return fmt.Errorf("failed to query retired CA cert: %w", err)
	} else if err == nil {
		if retiredCertExpiration.Before(time.Now()) {
			err = updateCertificateStatus(ctx, tx, retiredCertID, deleted)
			if err != nil {
				return fmt.Errorf("failed to update certificate status: %w", err)
			}
			log.Info("Updated retired cert to deleted", "certID", retiredCertID)
		}
	}

	if err = tx.Commit(); err != nil {
		return err
	}

	log.Info("Successfully reconciled certs")
	return nil
}

func (r *BannerReconciler) createCA(ctx context.Context, b *bannerAPI.Banner, caPoolID string, status string) error {
	log := r.Log.WithValues("banner", b.Name)

	// The below CA template is mimicking the CA that is created by emissary
	expiration := time.Now().AddDate(5, 0, 0) // Valid for 5 years
	template := &x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName: "emissary-ca",
		},
		NotBefore:             time.Now(),
		NotAfter:              expiration,
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{},
		IsCA:                  true,
		BasicConstraintsValid: true,
	}

	cert, privateKey, err := generateCert(template)
	if err != nil {
		return recerr.New(err, bannerAPI.PlatformSecretsCreationFailedReason)
	}

	// creating a private key ref for the cert, which will be stored in SM
	privateKeySecretName := "ca-private-key-" + b.Name
	version, err := r.storeInSecretManager(ctx, privateKey, privateKeySecretName)
	if err != nil {
		return fmt.Errorf("failed to store CA cert in Secret Manager: %w", err)
	}

	// creating a cert ref for the CA cert, which will be stored in SM
	certSecretName := "ca-cert-" + b.Name
	certVersion, err := r.storeInSecretManager(ctx, cert, certSecretName)
	if err != nil {
		return fmt.Errorf("failed to store CA cert in Secret Manager: %w", err)
	}

	log.Info("Successfully stored CA cert in Secret Manager", "certVersion", certVersion)

	// generate refs based off name + version
	privateKeyRef := privateKeySecretName + "-" + version
	certRef := certSecretName + "-" + certVersion

	// add record to database
	certID, err := r.addRecordToDatabase(ctx, caPoolID, certRef, privateKeyRef, status, expiration)
	if err != nil {
		return fmt.Errorf("failed to add CA cert to database: %w", err)
	}

	log.Info("Successfully created CA cert in database", "CA Cert ID", certID)
	return nil
}

// generates a cert and returns the cert and private key in base64 format for storage in SM
func generateCert(template *x509.Certificate) ([]byte, []byte, error) {
	// Create private key
	privKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate private key: %v", err)
	}

	// Generate the certificate
	caBytes, err := x509.CreateCertificate(rand.Reader, template, template, &privKey.PublicKey, privKey)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to create certificate: %v", err)
	}

	// Encode the certificate to PEM format
	caPEM := new(bytes.Buffer)
	err = pem.Encode(caPEM, &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: caBytes,
	})
	if err != nil {
		return nil, nil, fmt.Errorf("failed to encode certificate to PEM: %v", err)
	}

	// Encode the private key to PEM format
	caPrivKeyPEM := new(bytes.Buffer)
	privBytes, err := x509.MarshalECPrivateKey(privKey)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to marshal private key: %v", err)
	}
	err = pem.Encode(caPrivKeyPEM, &pem.Block{
		Type:  "EC PRIVATE KEY",
		Bytes: privBytes,
	})
	if err != nil {
		return nil, nil, fmt.Errorf("failed to encode private key to PEM: %v", err)
	}

	// convert to B64 for secrets -> the LS0tLS1CRUdJTiBSU0EgUFJ.... format
	caCertBase64 := base64.StdEncoding.EncodeToString(caPEM.Bytes())
	caPrivKeyBase64 := base64.StdEncoding.EncodeToString(caPrivKeyPEM.Bytes())

	return []byte(caCertBase64), []byte(caPrivKeyBase64), nil
}

func (r *BannerReconciler) addRecordToDatabase(ctx context.Context, caPoolID string, certRef string, privateKeyRef string, status string, expiration time.Time) (string, error) {
	tx, err := r.EdgeDB.BeginTx(ctx, &sql.TxOptions{})
	if err != nil {
		return "", err
	}

	defer func() {
		if err != nil {
			err = errors.Join(err, tx.Rollback())
		}
	}()

	var caCertID string
	err = tx.QueryRowContext(ctx, createCACert, caPoolID, status, certRef, privateKeyRef, expiration).Scan(&caCertID)
	if err != nil {
		return "", fmt.Errorf("failed to insert CA cert into database: %w", err)
	}

	if err = tx.Commit(); err != nil {
		return "", err
	}

	return caCertID, nil
}

func (r *BannerReconciler) storeInSecretManager(ctx context.Context, secretData []byte, name string) (string, error) {
	// Create a new SecretManagerClient
	smClient, err := r.SecretManager.NewWithOptions(ctx, r.ForemanProjectID)
	if err != nil {
		return "", fmt.Errorf("error creating secretmanager writer client, err: %v", err)
	}
	labels := map[string]string{
		secretMgrApi.SecretLabel:                  string(workloadApi.Platform),
		secretMgrApi.SecretTypeLabel:              "banner-ca",
		secretMgrApi.SecretOwnerLabel:             "edge",
		secretMgrApi.SecretNamespaceSelectorLabel: string(constants.PlatformNamespaceSelector),
	}
	err = smClient.AddSecret(ctx, name, secretData, labels, false, nil, "")
	if err != nil {
		return "", fmt.Errorf("error adding secret, secretID: %v, err: %v", name, err)
	}

	fullName, err := smClient.GetLatestSecretValueInfo(ctx, name)
	if err != nil {
		return "", fmt.Errorf("error getting secret version, secretID: %v, err: %v", name, err)
	}

	return filepath.Base(fullName.GetName()), nil
}

func updateCertificateStatus(ctx context.Context, tx *sql.Tx, certID string, status string) error {
	_, err := tx.ExecContext(ctx, "UPDATE ca_certificates SET status = $1 WHERE ca_cert_edge_id = $2", status, certID)
	if err != nil {
		return fmt.Errorf("failed to update certificate status: %w", err)
	}

	return nil
}

func (r *BannerReconciler) createEdgeIssuerDatabaseUser(b *bannerAPI.Banner) *sequelApi.DatabaseUser {
	hash := uuid.FromUUID(b.Status.ClusterInfraClusterEdgeID).Hash()
	edgeIssuerSAName := fmt.Sprintf("issuer-%s", hash)
	iamUsername := fmt.Sprintf("%s@%s.iam", edgeIssuerSAName, b.Spec.GCP.ProjectID)

	grant := sequelApi.Grant{
		Schema: "public",
		TableGrant: []sequelApi.TableGrant{
			{
				Table: "ca_pools",
				Permissions: []sequelApi.Permissions{
					{
						Permission: "SELECT",
					},
					{
						Permission: "INSERT",
					},
					{
						Permission: "UPDATE",
					},
				},
			},
			{
				Table: "ca_certificates",
				Permissions: []sequelApi.Permissions{
					{
						Permission: "SELECT",
					},
					{
						Permission: "INSERT",
					},
					{
						Permission: "UPDATE",
					},
				},
			},
		},
	}

	return &sequelApi.DatabaseUser{
		TypeMeta: gvkToTypeMeta(sequelApi.UserGVK),
		ObjectMeta: metav1.ObjectMeta{
			Name:      edgeIssuerSAName,
			Namespace: b.Name,
		},
		Spec: sequelApi.UserSpec{
			Type: sequelApi.CloudSAUserType,
			CommonOptions: sequelApi.CommonOptions{
				Prune: true,
				Force: true,
			},
			InstanceRef: sequelApi.InstanceReference{
				Name:      r.DatabaseName + dbInstance,
				ProjectID: r.ForemanProjectID,
			},
			ServiceAccount: &sequelApi.ServiceAccount{
				EmailRef:    fmt.Sprintf("%s.gserviceaccount.com", iamUsername),
				IAMUsername: iamUsername,
			},
			Grants: []sequelApi.Grant{
				grant,
			},
		},
	}
}