package main

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

	gohcl "github.com/hashicorp/hcl/v2/gohcl"
	hclparse "github.com/hashicorp/hcl/v2/hclparse"
	"github.com/spf13/cobra"
)

// CAConfig represents a CA block in HCL
// Example HCL:
//
//	ca "example_ca" {
//	    name = "My CA Name"
//	    country = "PL"
//	    organization = "ACME Corp."
//	}
type CAConfig struct {
	Label        string `hcl:",label"`
	Name         string `hcl:"name"`
	Country      string `hcl:"country"`
	Organization string `hcl:"organization"`
	SerialType   string `hcl:"serial_type"`
}

// ConfigFile represents the root of the HCL file
// It can contain multiple CA blocks
// Example: ca "example_ca" { ... }
type ConfigFile struct {
	CAs []CAConfig `hcl:"ca,block"`
}

// LoadCAConfig loads the CA configuration from an HCL file (first CA block)
func LoadCAConfig(path string) (*CAConfig, error) {
	parser := hclparse.NewParser()
	file, diags := parser.ParseHCLFile(path)
	if diags.HasErrors() {
		return nil, fmt.Errorf("failed to parse HCL: %s", diags.Error())
	}
	var configFile ConfigFile
	diags = gohcl.DecodeBody(file.Body, nil, &configFile)
	if diags.HasErrors() {
		return nil, fmt.Errorf("failed to decode HCL: %s", diags.Error())
	}
	if len(configFile.CAs) == 0 {
		return nil, fmt.Errorf("no 'ca' block found in config file")
	}
	ca := &configFile.CAs[0]
	if err := ca.Validate(); err != nil {
		return nil, err
	}
	return ca, nil
}

// GenerateCA generates a new CA certificate and private key
func GenerateCA(config *CAConfig) ([]byte, []byte, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return nil, nil, err
	}
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate serial number: %v", err)
	}
	tmpl := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Country:      []string{config.Country},
			Organization: []string{config.Organization},
			CommonName:   config.Name,
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(10, 0, 0),
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		BasicConstraintsValid: true,
		IsCA:                  true,
	}
	certDER, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, &priv.PublicKey, priv)
	if err != nil {
		return nil, nil, err
	}
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})
	return certPEM, keyPEM, nil
}

func SavePEM(filename string, data []byte, secure bool, overwrite bool) error {
	if !overwrite {
		if _, err := os.Stat(filename); err == nil {
			return fmt.Errorf("file %s already exists (overwrite not allowed)", filename)
		} else if !os.IsNotExist(err) {
			return fmt.Errorf("could not check file %s: %v", filename, err)
		}
	}

	if secure {
		return os.WriteFile(filename, data, 0600) // Read/write for owner only
	} else {
		return os.WriteFile(filename, data, 0644) // Read/write for owner, read for group and others
	}
}

// Validate checks required fields and sets defaults for CAConfig
func (c *CAConfig) Validate() error {
	if c.Name == "" {
		return fmt.Errorf("CA 'name' is required")
	}
	if c.Country == "" {
		return fmt.Errorf("CA 'country' is required")
	}
	if c.Organization == "" {
		return fmt.Errorf("CA 'organization' is required")
	}
	if c.SerialType == "" {
		c.SerialType = "random"
	}
	if c.SerialType != "random" && c.SerialType != "sequential" {
		return fmt.Errorf("CA 'serial_type' must be 'random' or 'sequential'")
	}
	return nil
}

func main() {
	var configPath string
	var overwrite bool

	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")

	rootCmd.AddCommand(initcaCmd)

	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 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()
	fmt.Println("Use 'lab-ca <command> --help' for more information about a command.")
}