package main

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

	"go.etcd.io/bbolt"
)

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

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

// --- Lego install ---

func runInstallLego() {
	arch, err := exec.Command("uname", "-m").Output()
	if err != nil {
		log.Fatalf("Failed to detect architecture: %v", err)
	}
	archStr := strings.TrimSpace(string(arch))
	archStr = strings.ReplaceAll(archStr, "x86_64", "amd64")
	archStr = strings.ReplaceAll(archStr, "aarch64", "arm64")

	// Get latest version
	out, err := exec.Command("curl", "-fsSL",
		"https://api.github.com/repos/go-acme/lego/releases/latest").Output()
	if err != nil {
		log.Fatalf("Failed to fetch latest lego version: %v", err)
	}
	version := extractLegoVersion(string(out))
	if version == "" {
		log.Fatalf("Could not parse latest lego version")
	}

	tarball := fmt.Sprintf("lego_%s_linux_%s.tar.gz", version, archStr)
	url := fmt.Sprintf("https://github.com/go-acme/lego/releases/download/%s/%s", version, tarball)

	log.Printf("Downloading lego %s for %s...", version, archStr)
	if err := exec.Command("curl", "-fLO", url).Run(); err != nil {
		log.Fatalf("Failed to download lego: %v", err)
	}

	log.Printf("Installing to /usr/local/bin/...")
	if err := exec.Command("sudo", "tar", "-xzf", tarball, "-C", "/usr/local/bin", "lego").Run(); err != nil {
		log.Fatalf("Failed to extract lego: %v", err)
	}

	os.Remove(tarball)
	log.Printf("[OK] lego %s installed at /usr/local/bin/lego", version)
}

func extractLegoVersion(body string) string {
	// Simple parser for the GitHub API response
	marker := `"tag_name": "`
	idx := strings.Index(body, marker)
	if idx < 0 {
		return ""
	}
	rest := body[idx+len(marker):]
	end := strings.Index(rest, `"`)
	if end < 0 {
		return ""
	}
	return rest[:end]
}

// --- 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/certificates", "Backup directory for certs")
	legoDir := fs.String("lego-dir", "/opt/backup", "Lego working directory")
	dryRun := fs.Bool("dry-run", false, "Skip LE, generate self-signed if no backup")
	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 LE certs in backup or lego cache
	needIssue := false
	for _, domain := range domains {
		certFile := findCertFile(domain, *backupDir, *legoDir)
		if certFile == "" {
			needIssue = true
			break
		}
		if !isCertFromLE(certFile) {
			log.Printf("[INFO] %s: cert is self-signed, reissuing", domain)
			needIssue = true
			break
		}
		if isCertExpired(certFile, 30*24*time.Hour) {
			log.Printf("[INFO] %s: cert expires soon, reissuing", domain)
			needIssue = true
			break
		}
	}

	if needIssue {
		if *dryRun {
			log.Printf("[INFO] Dry-run mode — generating self-signed fallback certs")
			generateSelfSigned(domains, *backupDir)
			return
		}
		// Try lego first
		legoPath, legoErr := exec.LookPath("lego")
		if legoErr == nil {
			log.Printf("[INFO] Requesting LE certificates via lego...")
			if err := obtainCertsLego(domains, *email, *backupDir, *legoDir, legoPath); err != nil {
				log.Printf("[WARN] lego failed: %v", err)
				log.Printf("[INFO] Generating self-signed fallback certs")
				generateSelfSigned(domains, *backupDir)
			}
		} else {
			log.Printf("[WARN] lego not found (install with 'nextwks-tool install-lego')")
			log.Printf("[INFO] Generating self-signed fallback certs")
			generateSelfSigned(domains, *backupDir)
		}
	} 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", domain, expiry.Format(time.RFC3339))
		}
	}
}

func obtainCertsLego(domains []string, email, backupDir, legoDir, legoPath string) error {
	if err := os.MkdirAll(legoDir, 0755); err != nil {
		return fmt.Errorf("creating lego dir: %w", err)
	}

	certDir := filepath.Join(legoDir, "certificates")

	for _, domain := range domains {
		log.Printf("  Requesting cert for %s...", domain)
		args := []string{"run",
			"--http",
			"--http.address=:80",
			"--path", legoDir,
			"--accept-tos",
			"-m", email,
			"-d", domain,
		}

		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		cmd := exec.CommandContext(ctx, legoPath, args...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if err := cmd.Run(); err != nil {
			log.Printf("  [WARN] lego failed for %s: %v", domain, err)
			// Don't try remaining domains — caller will generate self-signed fallbacks
			return fmt.Errorf("lego failed for %s: %w", domain, err)
		}

		// Small delay so port 80 is released between runs
		time.Sleep(2 * time.Second)

		// Copy cert to per-domain backup
		crtSrc := filepath.Join(certDir, domain+".crt")
		keySrc := filepath.Join(certDir, domain+".key")
		domainDir := filepath.Join(backupDir, domain)
		os.MkdirAll(domainDir, 0755)

		if fileExists(crtSrc) && fileExists(keySrc) {
			copyFile(crtSrc, filepath.Join(domainDir, "fullchain.pem"))
			copyFile(keySrc, filepath.Join(domainDir, "privkey.pem"))
			log.Printf("[OK] Certificate obtained 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 findCertFile(domain, backupDir, legoDir string) string {
	// Check backup first
	candidates := []string{
		filepath.Join(backupDir, domain, "fullchain.pem"),
		filepath.Join(legoDir, "certificates", domain+".crt"),
	}
	// Lego issues SAN cert named after first domain — check in lego cache
	entries, _ := os.ReadDir(filepath.Join(legoDir, "certificates"))
	for _, e := range entries {
		if strings.HasSuffix(e.Name(), ".crt") && !strings.Contains(e.Name(), ".issuer.") {
			candidates = append(candidates, filepath.Join(legoDir, "certificates", e.Name()))
			break
		}
	}
	for _, c := range candidates {
		if fileExists(c) && isCertFromLE(c) {
			return c
		}
	}
	return ""
}

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

func isCertFromLE(certFile string) bool {
	data, err := os.ReadFile(certFile)
	if err != nil {
		return false
	}
	block, _ := pem.Decode(data)
	if block == nil {
		return false
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return false
	}
	if len(cert.Issuer.Organization) > 0 {
		return strings.Contains(cert.Issuer.Organization[0], "Let's Encrypt")
	}
	return false
}

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

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

		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
}