// Copyright 2020-2026 Project Capsule Authors
// SPDX-License-Identifier: Apache-2.0

package cert

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"time"

	"github.com/pkg/errors"
)

type CA interface {
	GenerateCertificate(opts CertificateOptions) (certificatePem *bytes.Buffer, certificateKey *bytes.Buffer, err error)
	CACertificatePem() (b *bytes.Buffer, err error)
	CAPrivateKeyPem() (b *bytes.Buffer, err error)
	ExpiresIn(now time.Time) (time.Duration, error)
	ValidateCert(certificate *x509.Certificate) error
}

type CapsuleCA struct {
	certificate *x509.Certificate
	key         *rsa.PrivateKey
}

func NewCertificateAuthorityFromBytes(certBytes, keyBytes []byte) (*CapsuleCA, error) {
	cert, key, err := GetCertificateWithPrivateKeyFromBytes(certBytes, keyBytes)
	if err != nil {
		return nil, err
	}

	return &CapsuleCA{
		certificate: cert,
		key:         key,
	}, nil
}

func (c *CapsuleCA) ExpiresIn(now time.Time) (time.Duration, error) {
	if c == nil || c.certificate == nil {
		return 0, errors.New("CA certificate is nil")
	}

	return c.certificate.NotAfter.Sub(now), nil
}

func (c *CapsuleCA) ValidateCert(certificate *x509.Certificate) error {
	if c == nil || c.certificate == nil {
		return errors.New("CA certificate is nil")
	}

	if certificate == nil {
		return errors.New("certificate is nil")
	}

	roots := x509.NewCertPool()
	roots.AddCert(c.certificate)

	_, err := certificate.Verify(x509.VerifyOptions{
		Roots: roots,
	})

	return err
}

func (c CapsuleCA) CACertificatePem() (b *bytes.Buffer, err error) {
	var crtBytes []byte

	crtBytes, err = x509.CreateCertificate(rand.Reader, c.certificate, c.certificate, &c.key.PublicKey, c.key)
	if err != nil {
		return b, err
	}

	b = new(bytes.Buffer)
	err = pem.Encode(b, &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: crtBytes,
	})

	return b, err
}

func (c CapsuleCA) CAPrivateKeyPem() (b *bytes.Buffer, err error) {
	b = new(bytes.Buffer)

	return b, pem.Encode(b, &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(c.key),
	})
}

func ValidateCertificate(cert *x509.Certificate, key *rsa.PrivateKey, expirationThreshold time.Duration) error {
	if !key.PublicKey.Equal(cert.PublicKey) {
		return errors.New("certificate signed by wrong public key")
	}

	now := time.Now()

	if now.Before(cert.NotBefore) {
		return errors.New("certificate is not valid yet")
	}

	if now.After(cert.NotAfter.Add(-expirationThreshold)) {
		return errors.New("certificate expired or going to expire soon")
	}

	return nil
}

func GenerateCertificateAuthority() (s *CapsuleCA, err error) {
	s = &CapsuleCA{
		certificate: &x509.Certificate{
			SerialNumber: big.NewInt(2019),
			Subject: pkix.Name{
				Organization:  []string{"Projectcapsule"},
				Country:       []string{"UK"},
				Province:      []string{""},
				Locality:      []string{"London"},
				StreetAddress: []string{"27, Old Gloucester Street"},
				PostalCode:    []string{"WC1N 3AX"},
			},
			NotBefore:             time.Now(),
			NotAfter:              time.Now().AddDate(10, 0, 0),
			IsCA:                  true,
			ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
			KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
			BasicConstraintsValid: true,
		},
	}

	s.key, err = rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return nil, err
	}

	return s, err
}

func GetCertificateFromBytes(raw []byte) (*x509.Certificate, error) {
	block, _ := pem.Decode(raw)
	if block == nil {
		return nil, fmt.Errorf("failed to decode certificate PEM")
	}

	if block.Type != "CERTIFICATE" {
		return nil, fmt.Errorf("expected CERTIFICATE PEM block, got %q", block.Type)
	}

	certificate, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse certificate: %w", err)
	}

	return certificate, nil
}

func GetPrivateKeyFromBytes(raw []byte) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode(raw)
	if block == nil {
		return nil, fmt.Errorf("failed to decode private key PEM")
	}

	if block.Type != "RSA PRIVATE KEY" {
		return nil, fmt.Errorf("expected RSA PRIVATE KEY PEM block, got %q", block.Type)
	}

	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse RSA private key: %w", err)
	}

	return privateKey, nil
}

func GetCertificateWithPrivateKeyFromBytes(certBytes, keyBytes []byte) (*x509.Certificate, *rsa.PrivateKey, error) {
	cert, err := GetCertificateFromBytes(certBytes)
	if err != nil {
		return nil, nil, err
	}

	key, err := GetPrivateKeyFromBytes(keyBytes)
	if err != nil {
		return nil, nil, err
	}

	return cert, key, nil
}

func (c *CapsuleCA) GenerateCertificate(opts CertOpts) (certificatePem *bytes.Buffer, certificateKey *bytes.Buffer, err error) {
	var certPrivKey *rsa.PrivateKey

	certPrivKey, err = rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return nil, nil, err
	}

	sans := opts.SAN
	if sans.Empty() {
		return nil, nil, fmt.Errorf("cannot generate certificate without SANs")
	}

	cert := &x509.Certificate{
		SerialNumber: big.NewInt(1658),
		Subject: pkix.Name{
			Organization:  []string{"Projectcapsule"},
			Country:       []string{"UK"},
			Province:      []string{""},
			Locality:      []string{"London"},
			StreetAddress: []string{"27, Old Gloucester Street"},
			PostalCode:    []string{"WC1N 3AX"},
		},
		DNSNames:     sans.DNSNames,
		IPAddresses:  sans.IPAddrs,
		NotBefore:    time.Now().AddDate(0, 0, -1),
		NotAfter:     opts.GetExpirationDate(),
		SubjectKeyId: []byte{1, 2, 3, 4, 6},
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		KeyUsage:     x509.KeyUsageDigitalSignature,
	}

	var certBytes []byte

	certBytes, err = x509.CreateCertificate(rand.Reader, cert, c.certificate, &certPrivKey.PublicKey, c.key)
	if err != nil {
		return nil, nil, err
	}

	certificatePem = new(bytes.Buffer)

	err = pem.Encode(certificatePem, &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certBytes,
	})
	if err != nil {
		return certificatePem, certificateKey, err
	}

	certificateKey = new(bytes.Buffer)

	err = pem.Encode(certificateKey, &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(certPrivKey),
	})
	if err != nil {
		return certificatePem, certificateKey, err
	}

	return certificatePem, certificateKey, err
}
