package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"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"`
}

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"`
}

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"`
	Certificates []CertificateDefinition `hcl:"certificate,block"`
}

// Global CA configurationa and state variables
var CAState *_CAState
var CAConfig *_CAConfig

// LoadCA loads the CA config and sets the global CAConfig variable
func LoadCA(path string) error {
	parser := hclparse.NewParser()
	file, diags := parser.ParseHCLFile(path)
	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
	// Derive caStatePath from caConfig label and config file path
	caDir := filepath.Dir(path)
	caLabel := config.Current.Label
	caStatePath := filepath.Join(caDir, caLabel+"_state.json")
	err := error(nil)
	CAState, err = LoadCAState(caStatePath)
	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
}

func parseValidity(validity string) (time.Duration, error) {
	if validity == "" {
		return time.Hour * 24 * 365 * 5, nil // default 5 years
	}
	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 {
			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 GenerateCA() ([]byte, []byte, error) {
	// Use global CAConfig directly
	keySize := CAConfig.KeySize
	if keySize == 0 {
		keySize = 4096
	}
	priv, err := rsa.GenerateKey(rand.Reader, keySize)
	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)
	}
	validity, err := parseValidity(CAConfig.Validity)
	if err != nil {
		return nil, nil, err
	}
	now := time.Now()
	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 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)
	} 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(configPath string, overwrite bool) {
	if err := LoadCA(configPath); err != nil {
		fmt.Println("Error loading config:", err)
		return
	}
	// 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
		}
	}
	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
		}
	}
	certPEM, keyPEM, err := GenerateCA()
	if err != nil {
		fmt.Println("Error generating CA:", err)
		return
	}
	if err := SavePEM(filepath.Join(CAConfig.Paths.Certificates, "ca_cert.pem"), certPEM, false, overwrite); err != nil {
		fmt.Println("Error saving CA certificate:", err)
		return
	}
	if err := SavePEM(filepath.Join(CAConfig.Paths.PrivateKeys, "ca_key.pem"), keyPEM, true, overwrite); err != nil {
		fmt.Println("Error saving CA key:", err)
		return
	}
	fmt.Println("CA certificate and key generated.")
}

// Helper: issue a single certificate and key, save to files, return error if any
func issueSingleCertificate(def CertificateDefinition, overwrite, verbose bool) error {
	// Use global CAConfig directly
	// 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)
	}

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

	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
	if def.Validity != "" {
		validityDur, err = parseValidity(def.Validity)
		if err != nil {
			return fmt.Errorf("invalid validity value: %v", err)
		}
	} else {
		validityDur = 365 * 24 * time.Hour // default 1 year
	}

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

	certTmpl := x509.Certificate{
		SerialNumber: serialNumber,
		Subject:      subjectPKIX,
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(validityDur),
		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+"."+def.Type+".crt.pem")
	keyFile := filepath.Join(CAConfig.Paths.PrivateKeys, basename+"."+def.Type+".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,
		)
	}
	return nil
}

func IssueCertificate(configPath, subject, certType, validity string, san []string, name, fromFile string, overwrite, dryRun, verbose bool) {
	if fromFile != "" {
		certDefs, defaults, err := LoadCertificatesFile(fromFile)
		if err != nil {
			fmt.Printf("Error loading certificates file: %v\n", err)
			return
		}
		successes := 0
		errors := 0
		for i, def := range certDefs {
			if defaults != nil {
				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
				}
			}
			finalDef := renderCertificateDefTemplates(def, defaults)
			fmt.Printf("[%d/%d] Issuing %s... ", i+1, len(certDefs), finalDef.Name)
			if dryRun {
				fmt.Printf("(dry run)\n")
				successes++
				continue
			}
			err := issueSingleCertificate(finalDef, overwrite, verbose)
			if err != nil {
				fmt.Printf("error: %v\n", err)
				errors++
			} else {
				if !verbose {
					fmt.Printf("done\n")
				}
				successes++
			}
		}
		fmt.Printf("Batch complete: %d succeeded, %d failed.\n", successes, errors)
		// Save CA state after batch issuance
		caDir := filepath.Dir(configPath)
		caLabel := CAConfig.Label
		caStatePath := filepath.Join(caDir, caLabel+"_state.json")
		if err := SaveCAState(caStatePath, CAState); err != nil {
			fmt.Printf("Error saving CA state: %v\n", err)
		}
		return
	}
	// Single mode
	finalDef := renderCertificateDefTemplates(CertificateDefinition{Name: name, Subject: subject, Type: certType, Validity: validity, SAN: san}, nil)
	if dryRun {
		fmt.Printf("Would issue %s certificate for '%s' (dry run)\n", finalDef.Type, finalDef.Subject)
		return
	}
	err := issueSingleCertificate(finalDef, overwrite, verbose)
	if err != nil {
		fmt.Printf("Error: %v\n", err)
		return
	}
	fmt.Printf("%s certificate and key for '%s' generated.\n", finalDef.Type, finalDef.Subject)
	// Save CA state after single issuance
	caDir := filepath.Dir(configPath)
	caLabel := CAConfig.Label
	caStatePath := filepath.Join(caDir, caLabel+"_state.json")
	if err := SaveCAState(caStatePath, CAState); err != nil {
		fmt.Printf("Error saving CA state: %v\n", err)
	}
}

// Extract defaults from certificates.hcl (now using new LoadCertificatesFile signature)
func GetCertificateDefaults(path string) CertificateDefinition {
	_, defaults, err := LoadCertificatesFile(path)
	if err != nil || defaults == nil {
		return CertificateDefinition{}
	}
	return CertificateDefinition{
		Type:     defaults.Type,
		Validity: defaults.Validity,
		SAN:      defaults.SAN,
	}
}

// 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: apply Go template to a string using CertificateDefinition and CertificateDefaults as data
func applyTemplate(s string, def CertificateDefinition, defaults *CertificateDefaults) (string, error) {
	data := struct {
		CertificateDefinition
		Defaults *CertificateDefaults
	}{
		CertificateDefinition: def,
		Defaults:              defaults,
	}
	tmpl, err := template.New("").Parse(s)
	if err != nil {
		return s, err
	}
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, data); err != nil {
		return s, err
	}
	return buf.String(), nil
}

// Render all string fields in CertificateDefinition using Go templates and return a new CertificateDefinition
func renderCertificateDefTemplates(def CertificateDefinition, defaults *CertificateDefaults) CertificateDefinition {
	newDef := def
	// Subject: use def.Subject if set, else defaults.Subject (rendered)
	if def.Subject != "" {
		if rendered, err := applyTemplate(def.Subject, def, defaults); err == nil {
			newDef.Subject = rendered
		}
	} else if defaults != nil && defaults.Subject != "" {
		if rendered, err := applyTemplate(defaults.Subject, def, defaults); err == nil {
			newDef.Subject = rendered
		}
	}
	// Type: use def.Type if set, else defaults.Type (rendered)
	if def.Type != "" {
		if rendered, err := applyTemplate(def.Type, def, defaults); err == nil {
			newDef.Type = rendered
		}
	} else if defaults != nil && defaults.Type != "" {
		if rendered, err := applyTemplate(defaults.Type, def, defaults); err == nil {
			newDef.Type = rendered
		}
	}
	// Validity: use def.Validity if set, else defaults.Validity (rendered)
	if def.Validity != "" {
		if rendered, err := applyTemplate(def.Validity, def, defaults); err == nil {
			newDef.Validity = rendered
		}
	} else if defaults != nil && defaults.Validity != "" {
		if rendered, err := applyTemplate(defaults.Validity, def, defaults); err == nil {
			newDef.Validity = rendered
		}
	}
	// SAN: use def.SAN if set, else defaults.SAN (rendered)
	if len(def.SAN) > 0 {
		newSAN := make([]string, len(def.SAN))
		for i, s := range def.SAN {
			if rendered, err := applyTemplate(s, def, defaults); err == nil {
				newSAN[i] = rendered
			} else {
				newSAN[i] = s
			}
		}
		newDef.SAN = newSAN
	} else if defaults != nil && len(defaults.SAN) > 0 {
		newSAN := make([]string, len(defaults.SAN))
		for i, s := range defaults.SAN {
			if rendered, err := applyTemplate(s, def, defaults); err == nil {
				newSAN[i] = rendered
			} else {
				newSAN[i] = s
			}
		}
		newDef.SAN = newSAN
	}
	return newDef
}

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