package bannerctl

import (
	"context"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"math/big"
	"time"

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

	gcpProject "edge-infra.dev/pkg/lib/gcp/project"
	"edge-infra.dev/pkg/lib/uuid"

	ctrl "sigs.k8s.io/controller-runtime"

	bannerAPI "edge-infra.dev/pkg/edge/apis/banner/v1alpha1"
	"edge-infra.dev/pkg/edge/controllers/util/edgedb"
	"edge-infra.dev/pkg/edge/gcpinfra"
)

func (s *Suite) TestReconcileCAManagement() {
	bannerGUID := uuid.New().UUID
	generatedProjID := fmt.Sprintf("%s-%s", gcpinfra.ProjectIDPrefix, gcpProject.RandAN(29-(len(gcpinfra.ProjectIDPrefix))))
	banner := &bannerAPI.Banner{
		ObjectMeta: metav1.ObjectMeta{
			Name: bannerGUID,
		},
		Spec: bannerAPI.BannerSpec{
			DisplayName: "banner-for-ca-stuff",
			GCP: bannerAPI.GCPConfig{
				ProjectID: generatedProjID,
			},
			BSL: bannerAPI.BSLConfig{
				EnterpriseUnit: bannerAPI.BSLEnterpriseUnit{
					ID: uuid.New().UUID,
				},
				Organization: bannerAPI.BSLOrganization{
					Name: "test-org-for-ca-stuff",
				},
			},
		},
	}

	s.createBanner(banner)
	defer s.deleteBanner(banner)

	// Create a new BannerReconciler
	r := &BannerReconciler{
		Log:           ctrl.Log.WithName("controllers").WithName("Banner"),
		EdgeDB:        &edgedb.EdgeDB{DB: s.db},
		SecretManager: s.sMgr,
	}

	err := r.reconcileCerts(s.ctx, banner)
	s.Require().Nil(err)
}

func (s *Suite) TestCreateCAPool() {
	bannerGUID := uuid.New().UUID
	generatedProjID := fmt.Sprintf("%s-%s", gcpinfra.ProjectIDPrefix, gcpProject.RandAN(29-(len(gcpinfra.ProjectIDPrefix))))
	banner := &bannerAPI.Banner{
		ObjectMeta: metav1.ObjectMeta{
			Name: bannerGUID,
		},
		Spec: bannerAPI.BannerSpec{
			DisplayName: "banner-for-ca-stuff-2",
			GCP: bannerAPI.GCPConfig{
				ProjectID: generatedProjID,
			},
			BSL: bannerAPI.BSLConfig{
				EnterpriseUnit: bannerAPI.BSLEnterpriseUnit{
					ID: uuid.New().UUID,
				},
				Organization: bannerAPI.BSLOrganization{
					Name: "test-org-for-ca-stuff-2",
				},
			},
		},
	}

	s.createBanner(banner)
	defer s.deleteBanner(banner)

	// Create a new BannerReconciler
	r := &BannerReconciler{
		Log:           ctrl.Log.WithName("controllers").WithName("Banner"),
		EdgeDB:        &edgedb.EdgeDB{DB: s.db},
		SecretManager: s.sMgr,
	}

	poolID, err := r.reconcileCAPool(context.TODO(), bannerGUID)
	s.Require().Nil(err)
	s.Require().NotNil(poolID)

	// confirm returned poolID matches correct bannerID
	var bannerIDfromPool string
	err = s.db.QueryRow("SELECT banner_edge_id FROM ca_pools WHERE ca_pool_edge_id = $1", poolID).Scan(&bannerIDfromPool)
	s.Require().Nil(err)
	s.Require().Equal(bannerGUID, bannerIDfromPool)

	err = r.reconcileCACerts(s.ctx, banner, poolID)
	s.Require().Nil(err)

	// confirm the created ca cert using the ca pool's id
	var caCert string
	err = s.db.QueryRow("SELECT ca_cert_edge_id FROM ca_certificates WHERE ca_pool_edge_id = $1", poolID).Scan(&caCert)
	s.Require().Nil(err)
	s.Require().NotNil(caCert)
}

func (s *Suite) TestCreateCACert() {
	template := &x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName: "emissary-ca",
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(5, 0, 0),
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{},
		IsCA:                  true,
		BasicConstraintsValid: true,
	}

	cert, privateKey, err := generateCert(template)
	s.Require().Nil(err)
	s.Require().NotNil(cert)
	s.Require().NotNil(privateKey)

	// Decode the base64-encoded certificate
	_, err = base64.StdEncoding.Decode(privateKey, privateKey)
	s.Require().Nil(err)

	_, err = base64.StdEncoding.Decode(cert, cert)
	s.Require().Nil(err)

	// Decode the PEM-encoded certificate
	block, _ := pem.Decode(cert)
	s.Require().NotNil(block)
	s.Require().Equal("CERTIFICATE", block.Type)

	// Parse the certificate, confirm it's valid
	parsedCert, err := x509.ParseCertificate(block.Bytes)
	s.Require().Nil(err)
	s.Require().NotNil(parsedCert)

	// confirm the certificate is not expired
	s.Require().True(time.Now().Before(parsedCert.NotAfter))
	s.Require().True(parsedCert.IsCA)

	// Decode the PEM-encoded private key
	block, _ = pem.Decode(privateKey)
	s.Require().NotNil(block)
	s.Require().Equal("EC PRIVATE KEY", block.Type)

	// Parse the private key, confirm it's valid
	parsedPrivateKey, err := x509.ParseECPrivateKey(block.Bytes)
	s.Require().Nil(err)
	s.Require().NotNil(parsedPrivateKey)
}

func (s *Suite) TestReconcileCACert() {
	bannerGUID := uuid.New().UUID
	generatedProjID := fmt.Sprintf("%s-%s", gcpinfra.ProjectIDPrefix, gcpProject.RandAN(29-(len(gcpinfra.ProjectIDPrefix))))
	r := &BannerReconciler{
		Log:           ctrl.Log.WithName("controllers").WithName("Banner"),
		EdgeDB:        &edgedb.EdgeDB{DB: s.db},
		SecretManager: s.sMgr,
	}

	// Create a Banner instance
	banner := &bannerAPI.Banner{
		ObjectMeta: metav1.ObjectMeta{
			Name: bannerGUID,
		},
		Spec: bannerAPI.BannerSpec{
			DisplayName: "banner-for-ca-stuff-3",
			GCP: bannerAPI.GCPConfig{
				ProjectID: generatedProjID,
			},
			BSL: bannerAPI.BSLConfig{
				EnterpriseUnit: bannerAPI.BSLEnterpriseUnit{
					ID: uuid.New().UUID,
				},
				Organization: bannerAPI.BSLOrganization{
					Name: "test-org-for-ca-stuff-3",
				},
			},
		},
	}

	s.createBanner(banner)
	defer s.deleteBanner(banner)

	poolID, err := r.reconcileCAPool(context.TODO(), bannerGUID)
	s.Require().Nil(err)
	s.Require().NotNil(poolID)

	err = r.reconcileCACerts(s.ctx, banner, poolID)
	s.Require().Nil(err)
}

func (s *Suite) TestReconcileCACertsRotation() {
	bannerGUID := uuid.New().UUID
	generatedProjID := fmt.Sprintf("%s-%s", gcpinfra.ProjectIDPrefix, gcpProject.RandAN(29-(len(gcpinfra.ProjectIDPrefix))))
	banner := &bannerAPI.Banner{
		ObjectMeta: metav1.ObjectMeta{
			Name: bannerGUID,
		},
		Spec: bannerAPI.BannerSpec{
			DisplayName: "banner-for-ca-stuff-4",
			GCP: bannerAPI.GCPConfig{
				ProjectID: generatedProjID,
			},
			BSL: bannerAPI.BSLConfig{
				EnterpriseUnit: bannerAPI.BSLEnterpriseUnit{
					ID: uuid.New().UUID,
				},
				Organization: bannerAPI.BSLOrganization{
					Name: "test-org-for-ca-stuff-4",
				},
			},
		},
	}

	s.createBanner(banner)
	defer s.deleteBanner(banner)

	// Create a new BannerReconciler
	r := &BannerReconciler{
		Log:           ctrl.Log.WithName("controllers").WithName("Banner"),
		EdgeDB:        &edgedb.EdgeDB{DB: s.db},
		SecretManager: s.sMgr,
	}

	// Create CA pool
	poolID, err := r.reconcileCAPool(context.TODO(), bannerGUID)
	s.Require().Nil(err)
	s.Require().NotNil(poolID)

	// Insert active and staged CA certs
	activeCertID := uuid.New().UUID
	stagedCertID := uuid.New().UUID
	activeExpiration := time.Now().AddDate(0, 9, 0).Add(time.Second * 3) // 9 months from now
	stagedExpiration := time.Now().AddDate(5, 0, 0)                      // 5 years from now

	_, err = s.db.Exec(`
        INSERT INTO ca_certificates (ca_cert_edge_id, ca_pool_edge_id, status, cert_ref, private_key_ref, expiration)
        VALUES ($1, $2, 'active', 'active-cert-ref', 'active-private-key-ref', $3),
               ($4, $5, 'staged', 'staged-cert-ref', 'staged-private-key-ref', $6)
    `, activeCertID, poolID, activeExpiration, stagedCertID, poolID, stagedExpiration)
	s.Require().Nil(err)

	// Reconcile CA certs
	err = r.reconcileCACerts(s.ctx, banner, poolID)
	s.Require().Nil(err)

	// Verify that rotation does not occur before the rotation window is reached
	var status string
	err = s.db.QueryRow("SELECT status FROM ca_certificates WHERE ca_cert_edge_id = $1", activeCertID).Scan(&status)
	s.Require().Nil(err)
	s.Require().Equal("active", status)

	err = s.db.QueryRow("SELECT status FROM ca_certificates WHERE ca_cert_edge_id = $1", stagedCertID).Scan(&status)
	s.Require().Nil(err)
	s.Require().Equal("staged", status)

	// Here we will wait a few seconds for the active cert to reach the rotation window, call reconcileCACerts again
	// and verify that the certs are now rotated
	time.Sleep(time.Second * 4)

	// Reconcile CA certs
	err = r.reconcileCACerts(s.ctx, banner, poolID)
	s.Require().Nil(err)

	// Verify that the active cert is now retired and the staged cert is now active
	err = s.db.QueryRow("SELECT status FROM ca_certificates WHERE ca_cert_edge_id = $1", activeCertID).Scan(&status)
	s.Require().Nil(err)
	s.Require().Equal("retired", status)

	err = s.db.QueryRow("SELECT status FROM ca_certificates WHERE ca_cert_edge_id = $1", stagedCertID).Scan(&status)
	s.Require().Nil(err)
	s.Require().Equal("active", status)
}

func (s *Suite) TestReconcileCACertsCreateStagedCert() {
	bannerGUID := uuid.New().UUID
	generatedProjID := fmt.Sprintf("%s-%s", gcpinfra.ProjectIDPrefix, gcpProject.RandAN(29-(len(gcpinfra.ProjectIDPrefix))))
	banner := &bannerAPI.Banner{
		ObjectMeta: metav1.ObjectMeta{
			Name: bannerGUID,
		},
		Spec: bannerAPI.BannerSpec{
			DisplayName: "banner-for-ca-stuff-5",
			GCP: bannerAPI.GCPConfig{
				ProjectID: generatedProjID,
			},
			BSL: bannerAPI.BSLConfig{
				EnterpriseUnit: bannerAPI.BSLEnterpriseUnit{
					ID: uuid.New().UUID,
				},
				Organization: bannerAPI.BSLOrganization{
					Name: "test-org-for-ca-stuff-5",
				},
			},
		},
	}

	s.createBanner(banner)
	defer s.deleteBanner(banner)

	// Create a new BannerReconciler
	r := &BannerReconciler{
		Log:           ctrl.Log.WithName("controllers").WithName("Banner"),
		EdgeDB:        &edgedb.EdgeDB{DB: s.db},
		SecretManager: s.sMgr,
	}

	// Create CA pool
	poolID, err := r.reconcileCAPool(context.TODO(), bannerGUID)
	s.Require().Nil(err)
	s.Require().NotNil(poolID)

	// Insert an active CA cert
	activeCertID := uuid.New().UUID
	activeExpiration := time.Now().AddDate(0, 11, 0) // 11 months from now

	_, err = s.db.Exec(`
        INSERT INTO ca_certificates (ca_cert_edge_id, ca_pool_edge_id, status, cert_ref, private_key_ref, expiration)
        VALUES ($1, $2, 'active', 'active-cert-ref', 'active-private-key-ref', $3)
    `, activeCertID, poolID, activeExpiration)
	s.Require().Nil(err)

	// Reconcile CA certs
	err = r.reconcileCACerts(s.ctx, banner, poolID)
	s.Require().Nil(err)

	// Verify that a new staged cert has been created
	var stagedCertID string
	err = s.db.QueryRow("SELECT ca_cert_edge_id FROM ca_certificates WHERE ca_pool_edge_id = $1 AND status = 'staged'", poolID).Scan(&stagedCertID)
	s.Require().Nil(err)
	s.Require().NotNil(stagedCertID)

	// Verify that the active cert is still active
	var status string
	err = s.db.QueryRow("SELECT status FROM ca_certificates WHERE ca_cert_edge_id = $1", activeCertID).Scan(&status)
	s.Require().Nil(err)
	s.Require().Equal("active", status)
}

func (s *Suite) TestReconcileCACertsUpdateRetiredToDeleted() {
	bannerGUID := uuid.New().UUID
	generatedProjID := fmt.Sprintf("%s-%s", gcpinfra.ProjectIDPrefix, gcpProject.RandAN(29-(len(gcpinfra.ProjectIDPrefix))))
	banner := &bannerAPI.Banner{
		ObjectMeta: metav1.ObjectMeta{
			Name: bannerGUID,
		},
		Spec: bannerAPI.BannerSpec{
			DisplayName: "banner-for-ca-stuff-6",
			GCP: bannerAPI.GCPConfig{
				ProjectID: generatedProjID,
			},
			BSL: bannerAPI.BSLConfig{
				EnterpriseUnit: bannerAPI.BSLEnterpriseUnit{
					ID: uuid.New().UUID,
				},
				Organization: bannerAPI.BSLOrganization{
					Name: "test-org-for-ca-stuff-6",
				},
			},
		},
	}

	s.createBanner(banner)
	defer s.deleteBanner(banner)

	// Create a new BannerReconciler
	r := &BannerReconciler{
		Log:           ctrl.Log.WithName("controllers").WithName("Banner"),
		EdgeDB:        &edgedb.EdgeDB{DB: s.db},
		SecretManager: s.sMgr,
	}

	// Create CA pool
	poolID, err := r.reconcileCAPool(context.TODO(), bannerGUID)
	s.Require().Nil(err)
	s.Require().NotNil(poolID)

	// Insert a retired CA cert that is past its expiration
	retiredCertID := uuid.New().UUID
	expiredDate := time.Now().AddDate(-1, 0, 0) // 1 year ago
	activeCertID := uuid.New().UUID
	activeExpiration := time.Now().AddDate(5, 0, 0)

	_, err = s.db.Exec(`
        INSERT INTO ca_certificates (ca_cert_edge_id, ca_pool_edge_id, status, cert_ref, private_key_ref, expiration)
        VALUES 
            ($1, $2, 'retired', 'retired-cert-ref', 'retired-private-key-ref', $3),
            ($4, $5, 'active', 'active-cert-ref', 'active-private-key-ref', $6)
    `, retiredCertID, poolID, expiredDate, activeCertID, poolID, activeExpiration)
	s.Require().Nil(err)

	// Reconcile CA certs
	err = r.reconcileCACerts(s.ctx, banner, poolID)
	s.Require().Nil(err)

	// Verify that the retired cert is now deleted
	var status string
	err = s.db.QueryRow("SELECT status FROM ca_certificates WHERE ca_cert_edge_id = $1", retiredCertID).Scan(&status)
	s.Require().Nil(err)
	s.Require().Equal("deleted", status)
}