package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"text/template"
	"time"

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

type Paths struct {
	Certificates string `hcl:"certificates"`
	PrivateKeys  string `hcl:"private_keys"`
}

type _CAConfig struct {
	Label              string `hcl:",label"`
	Name               string `hcl:"name"`
	Country            string `hcl:"country"`
	Organization       string `hcl:"organization"`
	OrganizationalUnit string `hcl:"organizational_unit,optional"`
	Locality           string `hcl:"locality,optional"`
	Province           string `hcl:"province,optional"`
	Email              string `hcl:"email,optional"`
	SerialType         string `hcl:"serial_type,optional"`
	KeySize            int    `hcl:"key_size,optional"`
	Validity           string `hcl:"validity,optional"`
	Paths              Paths  `hcl:"paths,block"`
}

func (c *_CAConfig) StateName() string {
	return c.Label + "_state.json"
}

type Configuration struct {
	Current _CAConfig `hcl:"ca,block"`
}

type CertificateDefinition struct {
	Name     string   `hcl:",label"`
	Subject  string   `hcl:"subject,optional"`
	Type     string   `hcl:"type,optional"`
	Validity string   `hcl:"validity,optional"`
	SAN      []string `hcl:"san,optional"`
}

func (def *CertificateDefinition) fillDefaultValues(defaults *CertificateDefaults) {
	if defaults == nil {
		return
	}
	if def.Subject == "" {
		def.Subject = defaults.Subject
	}
	if def.Type == "" {
		def.Type = defaults.Type
	}
	if def.Validity == "" {
		def.Validity = defaults.Validity
	}
	if len(def.SAN) == 0 && len(defaults.SAN) > 0 {
		def.SAN = defaults.SAN
	}
}

// Helper: renderTemplates applies Go template to a string
// using the provided variables map. It returns an error if the template execution fails.
func applyTemplateToString(s string, variables map[string]string) (string, error) {
	tmpl, err := template.New("").Parse(s)
	if err != nil {
		return s, err
	}
	var buf bytes.Buffer
	err = tmpl.Execute(&buf, variables)
	if err != nil {
		return s, err
	}
	return buf.String(), nil
}

func (c *CertificateDefinition) RenderTemplates(variables map[string]string) error {
	// Apply Go templates to Subject and SAN fields using
	// the variables map
	if c.Subject != "" {
		renderedSubject, err := applyTemplateToString(c.Subject, variables)
		if err != nil {
			return fmt.Errorf("failed to render subject template: %v", err)
		}
		c.Subject = renderedSubject
	}

	if len(c.SAN) > 0 {
		for i, san := range c.SAN {
			renderedSAN, err := applyTemplateToString(san, variables)
			if err != nil {
				return fmt.Errorf("failed to render SAN template: %v", err)
			}
			c.SAN[i] = renderedSAN
		}
	}
	return nil
}

type CertificateDefaults struct {
	Subject  string   `hcl:"subject,optional"`
	Type     string   `hcl:"type,optional"`
	Validity string   `hcl:"validity,optional"`
	SAN      []string `hcl:"san,optional"`
}

type Certificates struct {
	Defaults     *CertificateDefaults    `hcl:"defaults,block"`
	Variables    map[string]string       `hcl:"variables,optional"`
	Certificates []CertificateDefinition `hcl:"certificate,block"`
}

// Load certificate provisioning configuration from the given path.
func (c *Certificates) LoadFromFile(path string) error {
	parser := hclparse.NewParser()
	file, diags := parser.ParseHCLFile(path)
	if diags.HasErrors() {
		return fmt.Errorf("failed to parse HCL: %s", diags.Error())
	}
	diags = gohcl.DecodeBody(file.Body, nil, c)
	if diags.HasErrors() {
		return fmt.Errorf("failed to decode HCL: %s", diags.Error())
	}

	return nil
}

// Global CA configuration and state variables
var CAConfigPath string
var CAState *_CAState
var CAConfig *_CAConfig
var CAKey *rsa.PrivateKey
var CACert *x509.Certificate

// LoadCAConfig parses and validates the CA config from the given path and stores it in the CAConfig global variable
func LoadCAConfig() error {
	fmt.Printf("Loading CA config from %s\n", CAConfigPath)
	parser := hclparse.NewParser()
	file, diags := parser.ParseHCLFile(CAConfigPath)
	if diags.HasErrors() {
		return fmt.Errorf("failed to parse HCL: %s", diags.Error())
	}
	var config Configuration
	diags = gohcl.DecodeBody(file.Body, nil, &config)
	if diags.HasErrors() {
		return fmt.Errorf("failed to decode HCL: %s", diags.Error())
	}
	if (_CAConfig{}) == config.Current {
		return fmt.Errorf("no 'ca' block found in config file")
	}
	if config.Current.Label == "" {
		return fmt.Errorf("the 'ca' block must have a label (e.g., ca \"mylabel\" {...})")
	}
	if err := config.Current.Validate(); err != nil {
		return err
	}
	CAConfig = &config.Current
	return nil
}

// LoadCA loads the CA config, certificate, key, and state
func LoadCA() error {
	var err error

	err = LoadCAConfig()
	if err != nil {
		return err
	}

	// Load CA key and certificate
	caCertPath := filepath.Join(CAConfig.Paths.Certificates, "ca_cert.pem")
	caKeyPath := filepath.Join(CAConfig.Paths.PrivateKeys, "ca_key.pem")

	caCertPEM, err := os.ReadFile(caCertPath)
	if err != nil {
		return fmt.Errorf("error reading CA certificate file: %v", err)
	}
	caKeyPEM, err := os.ReadFile(caKeyPath)
	if err != nil {
		return fmt.Errorf("error reading CA key file: %v", err)
	}

	caCertBlock, _ := pem.Decode(caCertPEM)
	if caCertBlock == nil {
		return fmt.Errorf("failed to parse CA certificate PEM")
	}
	CACert, err = x509.ParseCertificate(caCertBlock.Bytes)
	if err != nil {
		return fmt.Errorf("failed to parse CA certificate: %v", err)
	}
	caKeyBlock, _ := pem.Decode(caKeyPEM)
	if caKeyBlock == nil {
		return fmt.Errorf("failed to parse CA key PEM")
	}
	CAKey, err = x509.ParsePKCS1PrivateKey(caKeyBlock.Bytes)
	if err != nil {
		return fmt.Errorf("failed to parse CA private key: %v", err)
	}

	err = LoadCAState()
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("failed to load CA state: %w", err)
	}
	return nil
}

// Parse certificates.hcl file with defaults support
func LoadCertificatesFile(path string) ([]CertificateDefinition, *CertificateDefaults, error) {
	parser := hclparse.NewParser()
	file, diags := parser.ParseHCLFile(path)
	if diags.HasErrors() {
		return nil, nil, fmt.Errorf("failed to parse HCL: %s", diags.Error())
	}
	var certsFile Certificates
	diags = gohcl.DecodeBody(file.Body, nil, &certsFile)
	if diags.HasErrors() {
		return nil, nil, fmt.Errorf("failed to decode HCL: %s", diags.Error())
	}
	return certsFile.Certificates, certsFile.Defaults, nil
}

// Certificate definitions can have validity in various formats:
// - "1y" for 1 year
// - "6m" for 6 months
// - "30d" for 30 days
// Check the syntax and parse validity string into time.Duration
func parseValidity(validity string) (time.Duration, error) {
	// Return error is the function is called with an empty validity
	if validity == "" {
		return 0, fmt.Errorf("validity cannot be empty")
	}

	var n int
	var unit rune
	_, err := fmt.Sscanf(validity, "%d%c", &n, &unit)
	if err != nil {
		// If no unit, assume years
		_, err2 := fmt.Sscanf(validity, "%d", &n)
		if err2 != nil {
			// Still no success, return error
			return 0, fmt.Errorf("invalid validity format: %s", validity)
		}
		unit = 'y'
	}

	switch unit {
	case 'y':
		return time.Hour * 24 * 365 * time.Duration(n), nil
	case 'm':
		return time.Hour * 24 * 30 * time.Duration(n), nil
	case 'd':
		return time.Hour * 24 * time.Duration(n), nil
	default:
		return 0, fmt.Errorf("invalid validity unit: %c", unit)
	}
}

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)
	} else {
		return os.WriteFile(filename, data, 0644)
	}
}

func (p *Paths) Validate() error {
	if p.Certificates == "" {
		return fmt.Errorf("paths.certificates is required")
	}
	if p.PrivateKeys == "" {
		return fmt.Errorf("paths.private_keys is required")
	}
	return nil
}

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'")
	}
	if err := c.Paths.Validate(); err != nil {
		return err
	}
	return nil
}

func InitCA(overwrite bool) error {

	var err error

	err = LoadCAConfig()
	if err != nil {
		return err
	}

	// Create certificates directory with 0755, private keys with 0700
	if CAConfig.Paths.Certificates != "" {
		if err := os.MkdirAll(CAConfig.Paths.Certificates, 0755); err != nil {
			fmt.Printf("Error creating certificates directory '%s': %v\n", CAConfig.Paths.Certificates, err)
			return err
		}
	}
	if CAConfig.Paths.PrivateKeys != "" {
		if err := os.MkdirAll(CAConfig.Paths.PrivateKeys, 0700); err != nil {
			fmt.Printf("Error creating private keys directory '%s': %v\n", CAConfig.Paths.PrivateKeys, err)
			return err
		}
	}

	// Initialize CAState empty state with serial starting from 1
	CAState = &_CAState{
		Serial:       1, // Start serial from 1
		CreatedAt:    time.Now().UTC().Format(time.RFC3339),
		UpdatedAt:    time.Now().UTC().Format(time.RFC3339),
		Certificates: []CertificateRecord{},
	}

	keySize := CAConfig.KeySize
	if keySize == 0 {
		keySize = 4096
	}
	priv, err := rsa.GenerateKey(rand.Reader, keySize)
	if err != nil {
		return err
	}
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return fmt.Errorf("failed to generate serial number: %v", err)
	}

	if CAConfig.Validity == "" {
		CAConfig.Validity = "5y" // Use default validity of 5 years
	}

	validity, err := parseValidity(CAConfig.Validity)
	if err != nil {
		return err
	}

	now := time.Now()
	// Store CA certificate creation time
	CAState.CreatedAt = now.UTC().Format(time.RFC3339)

	tmpl := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Country:            []string{CAConfig.Country},
			Organization:       []string{CAConfig.Organization},
			OrganizationalUnit: optionalSlice(CAConfig.OrganizationalUnit),
			Locality:           optionalSlice(CAConfig.Locality),
			Province:           optionalSlice(CAConfig.Province),
			CommonName:         CAConfig.Name,
		},
		NotBefore:             now,
		NotAfter:              now.Add(validity),
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		BasicConstraintsValid: true,
		IsCA:                  true,
	}
	// Add email if present
	if CAConfig.Email != "" {
		tmpl.Subject.ExtraNames = append(tmpl.Subject.ExtraNames, pkix.AttributeTypeAndValue{
			Type:  []int{1, 2, 840, 113549, 1, 9, 1}, // emailAddress OID
			Value: CAConfig.Email,
		})
	}
	certDER, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, &priv.PublicKey, priv)
	if err != nil {
		return err
	}
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})

	if err := SavePEM(filepath.Join(CAConfig.Paths.Certificates, "ca_cert.pem"), certPEM, false, overwrite); err != nil {
		fmt.Println("Error saving CA certificate:", err)
		return err
	}
	if err := SavePEM(filepath.Join(CAConfig.Paths.PrivateKeys, "ca_key.pem"), keyPEM, true, overwrite); err != nil {
		fmt.Println("Error saving CA key:", err)
		return err
	}

	// set last updated time in the CAState
	CAState.UpdatedAt = time.Now().UTC().Format(time.RFC3339)
	// Save the state

	err = SaveCAState()
	if err != nil {
		fmt.Println("Error saving CA state:", err)
		return err
	}
	fmt.Println("CA certificate and key generated.")

	return nil
}

// Helper: issue a single certificate and key, save to files, return error if any
func issueSingleCertificate(def CertificateDefinition, overwrite, verbose bool) error {
	// Validate Name

	isValidName, err := regexp.MatchString(`^[A-Za-z0-9_-]+$`, def.Name)
	if err != nil {
		return fmt.Errorf("error validating certificate name: %v", err)
	}

	if !isValidName {
		return fmt.Errorf("certificate name must be specified and contain only letters, numbers, dash, or underscore")
	}

	// Initialize Subject if not specified
	if def.Subject == "" {
		def.Subject = def.Name
	}

	// Add default dns SAN for server/server-only if none specified
	if (def.Type == "server" || def.Type == "server-only") && len(def.SAN) == 0 {
		def.SAN = append(def.SAN, "dns:"+def.Subject)
	}

	priv, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return fmt.Errorf("failed to generate private key: %v", err)
	}
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return fmt.Errorf("failed to generate serial number: %v", err)
	}

	var validityDur time.Duration
	validity := def.Validity
	if validity == "" {
		validity = "1y" // default to 1 year
	}

	validityDur, err = parseValidity(validity)
	if err != nil {
		return fmt.Errorf("invalid validity value: %v", err)
	}

	var subjectPKIX pkix.Name
	if isDNFormat(def.Subject) {
		subjectPKIX = parseDistinguishedName(def.Subject)
	} else {
		subjectPKIX = pkix.Name{CommonName: def.Subject}
	}

	dateIssued := time.Now()
	expires := dateIssued.Add(validityDur)

	certTmpl := x509.Certificate{
		SerialNumber: serialNumber,
		Subject:      subjectPKIX,
		NotBefore:    dateIssued,
		NotAfter:     expires,
		KeyUsage:     x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
	}

	for _, s := range def.SAN {
		sLower := strings.ToLower(s)
		var val string
		if n, _ := fmt.Sscanf(sLower, "dns:%s", &val); n == 1 {
			certTmpl.DNSNames = append(certTmpl.DNSNames, val)
		} else if n, _ := fmt.Sscanf(sLower, "ip:%s", &val); n == 1 {
			certTmpl.IPAddresses = append(certTmpl.IPAddresses, net.ParseIP(val))
		} else if n, _ := fmt.Sscanf(sLower, "email:%s", &val); n == 1 {
			certTmpl.EmailAddresses = append(certTmpl.EmailAddresses, val)
		} else {
			return fmt.Errorf("invalid SAN format: %s", s)
		}
	}

	switch def.Type {
	case "client":
		certTmpl.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth}
	case "server":
		certTmpl.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth}
	case "server-only":
		certTmpl.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}
	case "code-signing":
		certTmpl.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageCodeSigning}
	case "email":
		certTmpl.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageEmailProtection}
	default:
		return fmt.Errorf("unknown certificate type. Use one of: client, server, server-only, code-signing, email")
	}

	certDER, err := x509.CreateCertificate(rand.Reader, &certTmpl, CACert, &priv.PublicKey, CAKey)
	if err != nil {
		return fmt.Errorf("failed to create certificate: %v", err)
	}
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})

	basename := def.Name
	if basename == "" {
		basename = def.Subject
	}
	certFile := filepath.Join(CAConfig.Paths.Certificates, basename+".crt.pem")
	keyFile := filepath.Join(CAConfig.Paths.PrivateKeys, basename+".key.pem")
	if err := SavePEM(certFile, certPEM, false, overwrite); err != nil {
		return fmt.Errorf("error saving certificate: %v", err)
	}
	if err := SavePEM(keyFile, keyPEM, true, overwrite); err != nil {
		return fmt.Errorf("error saving key: %v", err)
	}
	if verbose {
		fmt.Printf(`
Certificate:
  Name:     %s
  Subject:  %s
  Type:     %s
  Validity: %s
  SAN:      %v
`,
			def.Name,
			def.Subject,
			def.Type,
			def.Validity,
			def.SAN,
		)
	}
	CAState.UpdateCAStateAfterIssue(
		CAConfig.SerialType,
		basename,
		serialNumber,
		validityDur,
	)
	return nil
}

// A prototype of certificate provisioning function
func ProvisionCertificates(filePath string, overwrite bool, dryRun bool, verbose bool) error {
	err := LoadCA()

	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		os.Exit(1)
	}

	// Make an empty Certificates struct to hold the definitions
	certDefs := Certificates{}

	// Load certificates provisioning configuration from the file (HCL syntax)
	err = certDefs.LoadFromFile(filePath)
	if err != nil {
		return fmt.Errorf("Error loading certificates file: %v", err)
	}

	// The certificate provisioning file must contain at least one certificate definition
	if len(certDefs.Certificates) < 1 {
		return fmt.Errorf("No certificates defined in %s", filePath)
	}

	// We will be counting successes and errors
	successes := 0
	errors := 0

	// Loop through all certificate definitions
	// to render templates and fill missing fields from defaults
	for i := range certDefs.Certificates {
		// Fill missing fields from defaults, if provided
		certDefs.Certificates[i].fillDefaultValues(certDefs.Defaults)
		// Render templates in the definition using the variables map
		// with added definition name.
		variables := certDefs.Variables
		if variables == nil {
			variables = make(map[string]string)
		}
		variables["Name"] = certDefs.Certificates[i].Name
		err = certDefs.Certificates[i].RenderTemplates(variables)
		if err != nil {
			return fmt.Errorf("failed to render templates for certificate %s: %v", certDefs.Certificates[i].Name, err)
		}
	}

	// No errors so far, now we can issue certificates
	for i := range certDefs.Certificates {
		fmt.Printf("[%d/%d] Issuing %s... ", i+1, len(certDefs.Certificates), certDefs.Certificates[i].Name)

		if dryRun {
			fmt.Printf("(dry run)\n")
			successes++
			continue
		}

		err = issueSingleCertificate(certDefs.Certificates[i], overwrite, verbose)
		if err != nil {
			fmt.Printf("error: %v\n", err)
			errors++
		} else {
			if !verbose {
				fmt.Printf("done\n")
			}
			successes++
		}
	}

	fmt.Printf("Provisioning complete: %d succeeded, %d failed.\n", successes, errors)

	err = SaveCAState()
	if err != nil {
		fmt.Printf("Error saving CA state: %v\n", err)
	}

	return nil
}

func IssueCertificate(certDef CertificateDefinition, overwrite bool, dryRun bool, verbose bool) error {
	err := LoadCA()
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		os.Exit(1)
	}

	if certDef.Subject == "" {
		certDef.Subject = certDef.Name
	}

	// Render templates in the certificae subject and SAN fields
	variables := map[string]string{"Name": certDef.Name}
	certDef.RenderTemplates(variables)

	if dryRun {
		fmt.Printf("Would issue %s certificate for '%s' (dry run)\n", certDef.Type, certDef.Subject)
		return nil
	}

	err = issueSingleCertificate(certDef, overwrite, verbose)

	if err != nil {
		return err
	}

	fmt.Printf("%s certificate and key for '%s' generated.\n", certDef.Type, certDef.Subject)
	if err := SaveCAState(); err != nil {
		fmt.Printf("Error saving CA state: %v\n", err)
	}

	return nil
}

// Helper: check if string looks like a DN (contains at least CN=...)
func isDNFormat(s string) bool {
	return len(s) > 0 && strings.Contains(s, "CN=")
}

// Helper: parse DN string into pkix.Name (supports CN, C, O, OU, L, ST, emailAddress)
func parseDistinguishedName(dn string) pkix.Name {
	var name pkix.Name
	parts := strings.Split(dn, ",")
	for _, part := range parts {
		kv := strings.SplitN(strings.TrimSpace(part), "=", 2)
		if len(kv) != 2 {
			continue
		}
		key, val := strings.TrimSpace(kv[0]), strings.TrimSpace(kv[1])
		switch key {
		case "CN":
			name.CommonName = val
		case "C":
			name.Country = append(name.Country, val)
		case "O":
			name.Organization = append(name.Organization, val)
		case "OU":
			name.OrganizationalUnit = append(name.OrganizationalUnit, val)
		case "L":
			name.Locality = append(name.Locality, val)
		case "ST":
			name.Province = append(name.Province, val)
		case "emailAddress":
			name.ExtraNames = append(name.ExtraNames, pkix.AttributeTypeAndValue{
				Type:  []int{1, 2, 840, 113549, 1, 9, 1}, // emailAddress OID
				Value: val,
			})
		}
	}
	return name
}

// Helper: convert optional string to []string or nil
func optionalSlice(s string) []string {
	if s == "" {
		return nil
	}
	return []string{s}
}