package pki

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

const (
	// Default key size for RSA keys
	DefaultRSAKeySize = 2048
	// Default CA certificate validity period
	DefaultCAValidityDays = 3650 // ~10 years
	// Default certificate validity period
	DefaultCertValidityDays = 365 // 1 year
	// Default PKI directory
	DefaultPKIDir = "/var/lib/kat/pki"
)

// GenerateCA creates a new Certificate Authority key pair and certificate.
// It saves the private key and certificate to the specified paths.
func GenerateCA(pkiDir string, keyPath, certPath string) error {
	// Create PKI directory if it doesn't exist
	if err := os.MkdirAll(pkiDir, 0700); err != nil {
		return fmt.Errorf("failed to create PKI directory: %w", err)
	}

	// Generate RSA key
	key, err := rsa.GenerateKey(rand.Reader, DefaultRSAKeySize)
	if err != nil {
		return fmt.Errorf("failed to generate CA key: %w", err)
	}

	// Create self-signed certificate
	serialNumber, err := generateSerialNumber()
	if err != nil {
		return fmt.Errorf("failed to generate serial number: %w", err)
	}

	// Certificate template
	notBefore := time.Now()
	notAfter := notBefore.Add(time.Duration(DefaultCAValidityDays) * 24 * time.Hour)

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   "KAT Root CA",
			Organization: []string{"KAT System"},
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		BasicConstraintsValid: true,
		IsCA:                  true,
		MaxPathLen:            1, // Only allow one level of intermediate certs
	}

	// Create certificate
	derBytes, err := x509.CreateCertificate(
		rand.Reader,
		&template,
		&template, // Self-signed
		&key.PublicKey,
		key,
	)
	if err != nil {
		return fmt.Errorf("failed to create CA certificate: %w", err)
	}

	// Save private key
	keyOut, err := os.OpenFile(keyPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return fmt.Errorf("failed to open CA key file for writing: %w", err)
	}
	defer keyOut.Close()

	err = pem.Encode(keyOut, &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	})
	if err != nil {
		return fmt.Errorf("failed to write CA key to file: %w", err)
	}

	// Save certificate
	certOut, err := os.OpenFile(certPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return fmt.Errorf("failed to open CA certificate file for writing: %w", err)
	}
	defer certOut.Close()

	err = pem.Encode(certOut, &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: derBytes,
	})
	if err != nil {
		return fmt.Errorf("failed to write CA certificate to file: %w", err)
	}

	return nil
}

// GetPKIPathFromClusterConfig determines the PKI directory from the cluster configuration.
// If backupPath is provided, it uses the parent directory of backupPath.
// Otherwise, it uses the default PKI directory.
func GetPKIPathFromClusterConfig(backupPath string) string {
	if backupPath == "" {
		return DefaultPKIDir
	}
	
	// Use the parent directory of backupPath
	return filepath.Dir(backupPath) + "/pki"
}

// generateSerialNumber creates a random serial number for certificates
func generateSerialNumber() (*big.Int, error) {
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128) // 128 bits
	return rand.Int(rand.Reader, serialNumberLimit)
}

// LoadCACertificate loads a CA certificate from a file
func LoadCACertificate(certPath string) (*x509.Certificate, error) {
	certPEM, err := os.ReadFile(certPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read CA certificate file: %w", err)
	}

	block, _ := pem.Decode(certPEM)
	if block == nil || block.Type != "CERTIFICATE" {
		return nil, fmt.Errorf("failed to decode PEM block containing certificate")
	}

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

	return cert, nil
}

// LoadCAPrivateKey loads a CA private key from a file
func LoadCAPrivateKey(keyPath string) (*rsa.PrivateKey, error) {
	keyPEM, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read CA key file: %w", err)
	}

	block, _ := pem.Decode(keyPEM)
	if block == nil || block.Type != "RSA PRIVATE KEY" {
		return nil, fmt.Errorf("failed to decode PEM block containing private key")
	}

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

	return key, nil
}