package certmanager

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/safedep/dry/log"
	"golang.org/x/sync/singleflight"
)

const (
	maxSystemCABundleBytes int64         = 10 * 1024 * 1024
	certExpiryThreshold    time.Duration = 1 * time.Hour
	goosDarwin                           = "darwin"
	goosLinux                            = "linux"
	goosWindows                          = "windows"
)

// certManager implements the CertificateManager interface
type certManager struct {
	ca     *Certificate
	cache  CertificateCache
	config CertManagerConfig
	group  singleflight.Group
}

// NewCertificateManagerWithCA creates a new certificate manager with an existing CA certificate
func NewCertificateManagerWithCA(ca *Certificate, config CertManagerConfig) (CertificateManager, error) {
	if ca == nil {
		return nil, fmt.Errorf("CA certificate cannot be nil")
	}

	config.SetDefaults()
	if err := config.Validate(); err != nil {
		return nil, fmt.Errorf("invalid config: %w", err)
	}

	if ca.X509Cert == nil || ca.PrivKey == nil {
		parsedCA, err := parseCertificate(ca)
		if err != nil {
			return nil, fmt.Errorf("failed to parse CA certificate: %w", err)
		}

		ca = parsedCA
	}

	if ca.IsExpired(certExpiryThreshold) {
		return nil, fmt.Errorf("CA certificate is expired")
	}

	return &certManager{
		ca:     ca,
		cache:  NewInMemoryCache(),
		config: config,
	}, nil
}

// GetCA returns the Certificate Authority certificate
func (cm *certManager) GetCA() (*Certificate, error) {
	return cm.ca, nil
}

// GenerateCertForHost creates a certificate for the given hostname.
// Uses caching and singleflight to ensure only one goroutine generates
// a certificate for a given hostname at a time, preventing CPU starvation
// when many concurrent CONNECT requests arrive for the same host.
func (cm *certManager) GenerateCertForHost(hostname string) (*Certificate, error) {
	if cached, found := cm.cache.Get(hostname); found {
		if !cached.IsExpired(certExpiryThreshold) {
			return cached, nil
		}
	}

	result, err, _ := cm.group.Do(hostname, func() (interface{}, error) {
		// Re-check cache: another goroutine in a previous singleflight
		// group may have populated it while we were waiting.
		if cached, found := cm.cache.Get(hostname); found {
			if !cached.IsExpired(certExpiryThreshold) {
				return cached, nil
			}
		}

		cert, err := cm.generateHostCert(hostname)
		if err != nil {
			return nil, fmt.Errorf("failed to generate certificate for %s: %w", hostname, err)
		}

		cm.cache.Set(hostname, cert)
		return cert, nil
	})
	if err != nil {
		return nil, err
	}

	return result.(*Certificate), nil
}

// GetTLSConfig returns a tls.Config for the given hostname
func (cm *certManager) GetTLSConfig(hostname string) (*tls.Config, error) {
	cert, err := cm.GenerateCertForHost(hostname)
	if err != nil {
		return nil, err
	}

	tlsCert, err := tls.X509KeyPair(cert.Certificate, cert.PrivateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create tls.Certificate: %w", err)
	}

	return &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
		MinVersion:   tls.VersionTLS12,
	}, nil
}

func (cm *certManager) generateHostCert(hostname string) (*Certificate, error) {
	privKey, err := rsa.GenerateKey(rand.Reader, cm.config.KeySize)
	if err != nil {
		return nil, fmt.Errorf("failed to generate private key: %w", err)
	}

	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, fmt.Errorf("failed to generate serial number: %w", err)
	}

	notBefore := time.Now()
	notAfter := notBefore.Add(time.Duration(cm.config.HostCertValidityDays) * 24 * time.Hour)

	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName: hostname,
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		DNSNames:              []string{hostname},
	}

	caPrivKey, ok := cm.ca.PrivKey.(*rsa.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("CA private key is not RSA")
	}

	certDER, err := x509.CreateCertificate(rand.Reader, template, cm.ca.X509Cert, &privKey.PublicKey, caPrivKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create certificate: %w", err)
	}

	x509Cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, fmt.Errorf("failed to parse generated certificate: %w", err)
	}

	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certDER,
	})

	privKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privKey),
	})

	return &Certificate{
		Certificate: certPEM,
		PrivateKey:  privKeyPEM,
		X509Cert:    x509Cert,
		PrivKey:     privKey,
	}, nil
}

// GenerateCA generates a new self-signed CA certificate using the given configuration
func GenerateCA(config CertManagerConfig) (*Certificate, error) {
	config.SetDefaults()
	if err := config.Validate(); err != nil {
		return nil, fmt.Errorf("invalid config: %w", err)
	}

	privKey, err := rsa.GenerateKey(rand.Reader, config.KeySize)
	if err != nil {
		return nil, fmt.Errorf("failed to generate CA private key: %w", err)
	}

	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, fmt.Errorf("failed to generate serial number: %w", err)
	}

	notBefore := time.Now()
	notAfter := notBefore.Add(time.Duration(config.CAValidityDays) * 24 * time.Hour)

	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   "PMG Proxy CA",
			Organization: []string{"SafeDep PMG"},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		BasicConstraintsValid: true,
		IsCA:                  true,
	}

	certDER, err := x509.CreateCertificate(rand.Reader, template, template, &privKey.PublicKey, privKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create CA certificate: %w", err)
	}

	x509Cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, fmt.Errorf("failed to parse generated CA certificate: %w", err)
	}

	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certDER,
	})

	privKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privKey),
	})

	return &Certificate{
		Certificate: certPEM,
		PrivateKey:  privKeyPEM,
		X509Cert:    x509Cert,
		PrivKey:     privKey,
	}, nil
}

func parseCertificate(cert *Certificate) (*Certificate, error) {
	certBlock, _ := pem.Decode(cert.Certificate)
	if certBlock == nil {
		return nil, fmt.Errorf("failed to decode PEM certificate")
	}

	x509Cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse X509 certificate: %w", err)
	}

	keyBlock, _ := pem.Decode(cert.PrivateKey)
	if keyBlock == nil {
		return nil, fmt.Errorf("failed to decode PEM private key")
	}

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

	return &Certificate{
		Certificate: cert.Certificate,
		PrivateKey:  cert.PrivateKey,
		X509Cert:    x509Cert,
		PrivKey:     privKey,
	}, nil
}

// ParseTLSCertificate converts a Certificate to a tls.Certificate
// This is useful for integrating with libraries that expect tls.Certificate
func ParseTLSCertificate(cert *Certificate) (tls.Certificate, error) {
	tlsCert, err := tls.X509KeyPair(cert.Certificate, cert.PrivateKey)
	if err != nil {
		return tls.Certificate{}, fmt.Errorf("failed to create X509 key pair: %w", err)
	}

	// Populate Leaf field if we have the parsed X.509 certificate
	// This is important for libraries (like goproxy) that need to inspect the certificate
	if cert.X509Cert != nil {
		tlsCert.Leaf = cert.X509Cert
	}

	return tlsCert, nil
}

// GenerateCAWithSystemCA generates a self-signed certificate and appends system CA bundle
// content to the certificate bytes. If the system bundle is unavailable or too large to merge,
// it falls back to the PMG CA so proxy startup remains functional.
func GenerateCAWithSystemCA(config CertManagerConfig) (*Certificate, error) {
	caCert, err := GenerateCA(config)
	if err != nil {
		return nil, fmt.Errorf("failed to generate CA: %w", err)
	}

	caCertPEM := caCert.Certificate
	systemBundlePath := firstReadablePath(systemCABundleCandidates()...)

	// No system CA found. Continue using only PMG cert.
	if systemBundlePath == "" {
		log.Warnf("Skipping system CA bundle merge: No system CA bundle file found")
		return caCert, nil
	}

	info, err := os.Stat(systemBundlePath)
	if err != nil {
		return nil, fmt.Errorf("failed to stat system CA bundle %s: %w", systemBundlePath, err)
	}

	// We make sure there is a boundary on the size of CA bundle loaded
	// from the system. Beyond that, we just skip it and return PMG cert.
	if info.Size() > maxSystemCABundleBytes {
		log.Errorf(
			"Skipping system CA bundle merge: %s is too large (%d bytes > %d bytes)",
			systemBundlePath,
			info.Size(),
			maxSystemCABundleBytes,
		)

		return caCert, nil
	}

	systemBundle, err := os.ReadFile(systemBundlePath)
	if err != nil {
		return nil, fmt.Errorf("failed to read system CA bundle %s: %w", systemBundlePath, err)
	}

	caLen := int64(len(caCertPEM))
	sysLen := int64(len(systemBundle))
	const extra = int64(2)

	totalCap := caLen + sysLen + extra
	if totalCap > maxSystemCABundleBytes {
		log.Errorf(
			"Skipping system CA bundle merge: merged CA would be too large (%d bytes > %d bytes)",
			totalCap,
			maxSystemCABundleBytes,
		)

		return caCert, nil
	}

	merged := make([]byte, 0, int(totalCap))
	merged = append(merged, caCertPEM...)

	if len(merged) > 0 && merged[len(merged)-1] != '\n' {
		merged = append(merged, '\n')
	}
	merged = append(merged, systemBundle...)

	if len(merged) > 0 && merged[len(merged)-1] != '\n' {
		merged = append(merged, '\n')
	}

	return &Certificate{
		Certificate: merged,
		PrivateKey:  caCert.PrivateKey,
		X509Cert:    caCert.X509Cert,
		PrivKey:     caCert.PrivKey,
	}, nil
}

func firstReadablePath(paths ...string) string {
	for _, path := range paths {
		if path == "" {
			continue
		}

		info, err := os.Stat(path)
		if err != nil || info.IsDir() {
			continue
		}

		f, err := os.Open(path)
		if err != nil {
			continue
		}
		_ = f.Close()

		return path
	}

	return ""
}

func systemCABundleCandidates() []string {
	return systemCABundleCandidatesForOS(runtime.GOOS)
}

func systemCABundleCandidatesForOS(goos string) []string {
	var candidates []string
	appendIfSet := func(key string) {
		if value := os.Getenv(key); value != "" {
			candidates = append(candidates, value)
		}
	}

	appendIfSet("SSL_CERT_FILE")
	appendIfSet("CURL_CA_BUNDLE")

	switch goos {
	case goosDarwin:
		candidates = append(candidates,
			"/opt/homebrew/etc/openssl@3/cert.pem",
			"/usr/local/etc/openssl@3/cert.pem",
			"/etc/ssl/cert.pem",
		)
	case goosLinux:
		candidates = append(candidates,
			"/etc/ssl/certs/ca-certificates.crt",
			"/etc/pki/tls/certs/ca-bundle.crt",
			"/etc/pki/ca-trust/extracted/pem/tls-ca-bundle.pem",
			"/etc/ssl/ca-bundle.pem",
			"/etc/ssl/cert.pem",
		)
	case goosWindows:
		programFiles := os.Getenv("ProgramFiles")
		programFilesX86 := os.Getenv("ProgramFiles(x86)")
		systemRoot := os.Getenv("SystemRoot")

		if programFiles != "" {
			candidates = append(candidates,
				filepath.Join(programFiles, "Git", "mingw64", "ssl", "certs", "ca-bundle.crt"),
				filepath.Join(programFiles, "Git", "usr", "ssl", "certs", "ca-bundle.crt"),
			)
		}
		if programFilesX86 != "" {
			candidates = append(candidates,
				filepath.Join(programFilesX86, "Git", "mingw32", "ssl", "certs", "ca-bundle.crt"),
			)
		}
		if systemRoot != "" {
			candidates = append(candidates,
				filepath.Join(systemRoot, "System32", "curl-ca-bundle.crt"),
			)
		}
	}

	return candidates
}