package main

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

	"github.com/spf13/cobra"
)

func main() {
	var configPath string
	var overwrite bool
	var subject string

	var rootCmd = &cobra.Command{
		Use:   "lab-ca",
		Short: "Certificate Authority Utility",
		Long:  "lab-ca - Certificate Authority Utility",
		Run: func(cmd *cobra.Command, args []string) {
			printMainHelp()
		},
	}

	var initcaCmd = &cobra.Command{
		Use:   "initca",
		Short: "Generate a new CA certificate and key",
		Run: func(cmd *cobra.Command, args []string) {
			handleInitCA(configPath, overwrite)
		},
	}

	initcaCmd.Flags().StringVar(&configPath, "config", "ca_config.hcl", "Path to CA configuration file")
	initcaCmd.Flags().BoolVar(&overwrite, "overwrite", false, "Allow overwriting existing files")

	var issueCmd = &cobra.Command{
		Use:   "issue",
		Short: "Issue a new client/server certificate",
		Run: func(cmd *cobra.Command, args []string) {
			handleIssue(configPath, subject, overwrite)
		},
	}

	issueCmd.Flags().StringVar(&configPath, "config", "ca_config.hcl", "Path to CA configuration file")
	issueCmd.Flags().StringVar(&subject, "subject", "", "Subject Common Name for the certificate (required)")
	issueCmd.Flags().BoolVar(&overwrite, "overwrite", false, "Allow overwriting existing files")
	issueCmd.MarkFlagRequired("subject")

	rootCmd.AddCommand(initcaCmd)
	rootCmd.AddCommand(issueCmd)

	if err := rootCmd.Execute(); err != nil {
		os.Exit(1)
	}
}

func handleInitCA(configPath string, overwrite bool) {
	config, err := LoadCAConfig(configPath)
	if err != nil {
		fmt.Println("Error loading config:", err)
		return
	}
	certPEM, keyPEM, err := GenerateCA(config)
	if err != nil {
		fmt.Println("Error generating CA:", err)
		return
	}
	if err := SavePEM("ca_cert.pem", certPEM, false, overwrite); err != nil {
		fmt.Println("Error saving CA certificate:", err)
		return
	}
	if err := SavePEM("ca_key.pem", keyPEM, true, overwrite); err != nil {
		fmt.Println("Error saving CA key:", err)
		return
	}
	fmt.Println("CA certificate and key generated.")
}

func handleIssue(configPath, subject string, overwrite bool) {
	// Load existing CA certificate and key from files
	caCertPEM, err := os.ReadFile("ca_cert.pem")
	if err != nil {
		fmt.Println("Error reading CA certificate file:", err)
		return
	}
	caKeyPEM, err := os.ReadFile("ca_key.pem")
	if err != nil {
		fmt.Println("Error reading CA key file:", err)
		return
	}

	// Parse CA cert and key
	caCertBlock, _ := pem.Decode(caCertPEM)
	if caCertBlock == nil {
		fmt.Println("Failed to parse CA certificate PEM")
		return
	}
	caCert, err := x509.ParseCertificate(caCertBlock.Bytes)
	if err != nil {
		fmt.Println("Failed to parse CA certificate:", err)
		return
	}
	caKeyBlock, _ := pem.Decode(caKeyPEM)
	if caKeyBlock == nil {
		fmt.Println("Failed to parse CA key PEM")
		return
	}
	caKey, err := x509.ParsePKCS1PrivateKey(caKeyBlock.Bytes)
	if err != nil {
		fmt.Println("Failed to parse CA private key:", err)
		return
	}

	// Generate client/server key
	priv, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		fmt.Println("Failed to generate private key:", err)
		return
	}
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		fmt.Println("Failed to generate serial number:", err)
		return
	}
	certTmpl := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName: subject,
		},
		NotBefore:   time.Now(),
		NotAfter:    time.Now().AddDate(1, 0, 0), // 1 year
		KeyUsage:    x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
	}
	certDER, err := x509.CreateCertificate(rand.Reader, &certTmpl, caCert, &priv.PublicKey, caKey)
	if err != nil {
		fmt.Println("Failed to create certificate:", err)
		return
	}
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})

	certFile := subject + ".crt.pem"
	keyFile := subject + ".key.pem"
	if err := SavePEM(certFile, certPEM, false, overwrite); err != nil {
		fmt.Println("Error saving certificate:", err)
		return
	}
	if err := SavePEM(keyFile, keyPEM, true, overwrite); err != nil {
		fmt.Println("Error saving key:", err)
		return
	}
	fmt.Printf("Certificate and key for '%s' generated.\n", subject)
}

func printMainHelp() {
	fmt.Println("lab-ca - Certificate Authority Utility")
	fmt.Println()
	fmt.Println("Usage:")
	fmt.Println("  lab-ca <command> [options]")
	fmt.Println()
	fmt.Println("Available commands:")
	fmt.Println("  initca     Generate a new CA certificate and key")
	fmt.Println("  issue      Issue a new client/server certificate")
	fmt.Println()
	fmt.Println("Use 'lab-ca <command> --help' for more information about a command.")
}