package proxy

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"math/big"
	"time"

	"git.secnex.io/secnex/masterlog"
	"git.secnex.io/secnex/pgproxy/config"
)

// generateSelfSignedCert generates a self-signed certificate for the given hostnames
func generateSelfSignedCert(hostnames []string) (tls.Certificate, error) {
	// Generate private key
	privKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return tls.Certificate{}, err
	}

	// Create certificate template
	notBefore := time.Now()
	notAfter := notBefore.Add(365 * 24 * time.Hour * 10) // Valid for 10 years

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return tls.Certificate{}, err
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization:  []string{"PostgreSQL Proxy"},
			Country:       []string{"DE"},
			Province:      []string{""},
			Locality:      []string{""},
			StreetAddress: []string{""},
			PostalCode:    []string{""},
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	// Set CommonName to first hostname
	if len(hostnames) > 0 {
		template.Subject.CommonName = hostnames[0]
	}

	// Add all hostnames as DNS names (SAN)
	template.DNSNames = hostnames

	// Create certificate
	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &privKey.PublicKey, privKey)
	if err != nil {
		return tls.Certificate{}, err
	}

	// Parse certificate
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return tls.Certificate{}, err
	}

	// Create TLS certificate
	return tls.Certificate{
		Certificate: [][]byte{certDER},
		PrivateKey:  privKey,
		Leaf:        cert,
	}, nil
}

// getCertificateForMappings generates or loads a certificate for the given mappings
func getCertificateForMappings(config *config.Config) (tls.Certificate, error) {
	// If certificate files are specified, try to load them
	if config.TLS.CertFile != "" && config.TLS.KeyFile != "" {
		cert, err := tls.LoadX509KeyPair(config.TLS.CertFile, config.TLS.KeyFile)
		if err == nil {
			masterlog.Info("Loaded TLS certificate from files", map[string]interface{}{
				"certFile": config.TLS.CertFile,
				"keyFile":  config.TLS.KeyFile,
			})
			return cert, nil
		}
		// If loading failed, log warning and fall through to generate
		masterlog.Info("Failed to load TLS certificate, generating self-signed certificate", map[string]interface{}{
			"error":    err,
			"certFile": config.TLS.CertFile,
			"keyFile":  config.TLS.KeyFile,
		})
	}

	// Generate self-signed certificate for all external hostnames
	hostnames := make([]string, 0, len(config.Mappings))
	for _, mapping := range config.Mappings {
		hostnames = append(hostnames, mapping.External)
	}

	// If no hostnames, use a default
	if len(hostnames) == 0 {
		hostnames = []string{"localhost"}
	}

	masterlog.Info("Generating self-signed certificate", map[string]interface{}{
		"hostnames": hostnames,
	})

	cert, err := generateSelfSignedCert(hostnames)
	if err != nil {
		return tls.Certificate{}, err
	}

	masterlog.Info("Generated self-signed certificate", map[string]interface{}{
		"hostnames":  hostnames,
		"validUntil": cert.Leaf.NotAfter.Format(time.RFC3339),
	})

	return cert, nil
}