package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"flag"
	"fmt"
	"log"
	"math/big"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"go.etcd.io/bbolt"
	"golang.org/x/crypto/acme/autocert"
)

func main() {
	if len(os.Args) < 2 {
		log.Fatalf("Usage: %s <cert|db> [flags]", os.Args[0])
	}

	switch os.Args[1] {
	case "cert":
		runCert(os.Args[2:])
	case "db":
		runDB(os.Args[2:])
	default:
		log.Fatalf("Unknown command: %s (use cert or db)", os.Args[1])
	}
}

// --- Cert command ---

func runCert(args []string) {
	fs := flag.NewFlagSet("cert", flag.ExitOnError)
	domainsStr := fs.String("domains", "", "Comma-separated domain list")
	email := fs.String("email", "", "ACME email")
	backupDir := fs.String("backup-dir", "/opt/backup/certs", "Backup directory for certs")
	deployDir := fs.String("deploy-dir", "", "Optional deploy directory to copy certs to")
	fs.Parse(args)

	if *domainsStr == "" || *email == "" {
		log.Fatal("--domains and --email are required")
	}

	domains := strings.Split(*domainsStr, ",")
	for i := range domains {
		domains[i] = strings.TrimSpace(domains[i])
	}

	if err := os.MkdirAll(*backupDir, 0755); err != nil {
		log.Fatalf("Failed to create backup dir: %v", err)
	}

	// Check if all domains have valid certs in backup
	needIssue := false
	for _, domain := range domains {
		certFile := filepath.Join(*backupDir, domain, "fullchain.pem")
		keyFile := filepath.Join(*backupDir, domain, "privkey.pem")
		if !fileExists(certFile) || !fileExists(keyFile) {
			needIssue = true
			break
		}
		// Check expiry
		if isCertExpired(certFile, 7*24*time.Hour) {
			log.Printf("[INFO] Cert for %s expires soon or is invalid, reissuing", domain)
			needIssue = true
			break
		}
	}

	if needIssue {
		log.Printf("[INFO] Requesting LE certificates for %v...", domains)
		if err := obtainCerts(domains, *email, *backupDir); err != nil {
			log.Printf("[WARN] LE cert issuance failed: %v", err)
			log.Printf("[INFO] Generating self-signed fallback certs")
			if err := generateSelfSigned(domains, *backupDir); err != nil {
				log.Printf("[WARN] Self-signed fallback also failed: %v", err)
			}
		}
	} else {
		log.Printf("[OK] All certificates found in backup (dry-run)")
		for _, domain := range domains {
			certFile := filepath.Join(*backupDir, domain, "fullchain.pem")
			expiry := getCertExpiry(certFile)
			log.Printf("  %s — expires %s (dry-run, deploy.sh copies from backup)", domain, expiry.Format(time.RFC3339))
		}
	}

	// Copy to deploy-dir if specified (legacy, not used by deploy.sh)
	if *deployDir != "" {
		if err := os.MkdirAll(*deployDir, 0755); err != nil {
			log.Printf("[WARN] Failed to create deploy dir: %v", err)
			return
		}
		for _, domain := range domains {
			srcCert := filepath.Join(*backupDir, domain, "fullchain.pem")
			srcKey := filepath.Join(*backupDir, domain, "privkey.pem")
			dstCert := filepath.Join(*deployDir, domain+".crt")
			dstKey := filepath.Join(*deployDir, domain+".key")

			if fileExists(srcCert) && fileExists(srcKey) {
				copyFile(srcCert, dstCert)
				copyFile(srcKey, dstKey)
				log.Printf("[OK] Deployed cert for %s", domain)
			}
		}
	}
}

func obtainCerts(domains []string, email, backupDir string) error {
	m := &autocert.Manager{
		Prompt: autocert.AcceptTOS,
		Email:  email,
		Cache:  autocert.DirCache(backupDir),
	}

	// Try to obtain certs by starting a temporary HTTP server for challenge
	ln, err := net.Listen("tcp", ":80")
	if err != nil {
		// Port 80 busy — try to use the autocert client directly without HTTP server
		log.Printf("[WARN] Port 80 not available (%v), trying direct ACME...", err)
		return obtainCertsDirect(domains, email, backupDir)
	}
	defer ln.Close()

	// Serve HTTP-01 challenge handler
	srv := &http.Server{
		Handler: m.HTTPHandler(nil),
		Addr:    ":80",
	}
	go srv.Serve(ln)

	// Give LE a moment to validate
	time.Sleep(2 * time.Second)

	for _, domain := range domains {
		hello := &tls.ClientHelloInfo{
			ServerName: domain,
		}
		cert, err := m.GetCertificate(hello)
		if err != nil {
			log.Printf("[WARN] Failed to get cert for %s: %v", domain, err)
			continue
		}
		// Store the cert to backup
		domainDir := filepath.Join(backupDir, domain)
		os.MkdirAll(domainDir, 0755)

		for _, c := range cert.Certificate {
			block := &pem.Block{Type: "CERTIFICATE", Bytes: c}
			f, err := os.OpenFile(filepath.Join(domainDir, "fullchain.pem"),
				os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
			if err != nil {
				return fmt.Errorf("writing cert: %w", err)
			}
			pem.Encode(f, block)
			f.Close()
		}

		// Extract and save private key
		if key, ok := cert.PrivateKey.(*rsa.PrivateKey); ok {
			keyBlock := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)}
			os.WriteFile(filepath.Join(domainDir, "privkey.pem"),
				pem.EncodeToMemory(keyBlock), 0600)
		}
		log.Printf("[OK] Certificate obtained for %s", domain)
	}

	return nil
}

func obtainCertsDirect(domains []string, email, backupDir string) error {
	// Direct ACME without port 80 — will likely fail but try anyway
	// This is a simplified fallback
	return fmt.Errorf("port 80 required for HTTP-01 challenge")
}

func generateSelfSigned(domains []string, backupDir string) error {
	for _, domain := range domains {
		key, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return err
		}

		serial, _ := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
		tmpl := &x509.Certificate{
			SerialNumber: serial,
			Subject:      pkix.Name{CommonName: domain},
			NotBefore:    time.Now(),
			NotAfter:     time.Now().Add(365 * 24 * time.Hour),
			KeyUsage:     x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
			ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		}
		if len(domains) > 1 {
			tmpl.DNSNames = domains
		}

		certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
		if err != nil {
			return err
		}

		domainDir := filepath.Join(backupDir, domain)
		os.MkdirAll(domainDir, 0755)

		certFile := filepath.Join(domainDir, "fullchain.pem")
		keyFile := filepath.Join(domainDir, "privkey.pem")

		f, _ := os.Create(certFile)
		pem.Encode(f, &pem.Block{Type: "CERTIFICATE", Bytes: certDER})
		f.Close()

		f, _ = os.Create(keyFile)
		pem.Encode(f, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
		f.Close()

		log.Printf("[INFO] Self-signed cert generated for %s", domain)
	}
	return nil
}

// --- DB command ---

func runDB(args []string) {
	fs := flag.NewFlagSet("db", flag.ExitOnError)
	dbPath := fs.String("db", "", "Path to BoltDB file")
	set := fs.String("set", "", "bucket:key:json-value")
	fs.Parse(args)

	if *dbPath == "" || *set == "" {
		log.Fatal("--db and --set are required")
	}

	parts := strings.SplitN(*set, ":", 3)
	if len(parts) != 3 {
		log.Fatalf("Invalid --set format. Use bucket:key:json-value")
	}
	bucket := parts[0]
	key := parts[1]
	value := parts[2]

	db, err := bbolt.Open(*dbPath, 0600, &bbolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		log.Fatalf("Failed to open BoltDB: %v", err)
	}
	defer db.Close()

	if err := db.Update(func(tx *bbolt.Tx) error {
		b, err := tx.CreateBucketIfNotExists([]byte(bucket))
		if err != nil {
			return err
		}
		return b.Put([]byte(key), []byte(value))
	}); err != nil {
		log.Fatalf("Failed to write to BoltDB: %v", err)
	}

	log.Printf("[OK] Wrote %s:%s to %s", bucket, key, *dbPath)
}

// --- Helpers ---

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func isCertExpired(certFile string, threshold time.Duration) bool {
	data, err := os.ReadFile(certFile)
	if err != nil {
		return true
	}
	block, _ := pem.Decode(data)
	if block == nil {
		return true
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return true
	}
	return time.Now().Add(threshold).After(cert.NotAfter)
}

func getCertExpiry(certFile string) time.Time {
	data, err := os.ReadFile(certFile)
	if err != nil {
		return time.Time{}
	}
	block, _ := pem.Decode(data)
	if block == nil {
		return time.Time{}
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return time.Time{}
	}
	return cert.NotAfter
}

func copyFile(src, dst string) error {
	data, err := os.ReadFile(src)
	if err != nil {
		return err
	}
	return os.WriteFile(dst, data, 0644)
}