package main

import (
	"bytes"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"math/big"
	"os"
	"text/template"
	"time"
)

// CertConfig holds paths and settings for certificate and kubeconfig generation.
type CertConfig struct {
	ServerURL    string        // written as the cluster.server in the returned kubeconfig
	ServerCACert string        // path — embedded in the kubeconfig so kubectl can verify the API server
	ClientCACert string        // path — signs the per-user client certificates
	ClientCAKey  string        // path
	Duration     time.Duration // validity period for generated client certs
	ClusterName  string        // name used for the cluster/context in generated kubeconfigs
}

// KubeconfigGenerator loads the k3s CAs once and issues per-user kubeconfigs on demand.
type KubeconfigGenerator struct {
	cfg      *CertConfig
	caCert   *x509.Certificate
	caKey    crypto.PrivateKey
	serverCA string // base64-encoded PEM of the server CA, ready to paste into kubeconfig
}

// NewKubeconfigGenerator reads the k3s CA files and prepares the generator.
func NewKubeconfigGenerator(cfg *CertConfig) (*KubeconfigGenerator, error) {
	caCert, caKey, err := loadCA(cfg.ClientCACert, cfg.ClientCAKey)
	if err != nil {
		return nil, fmt.Errorf("loading client CA: %w", err)
	}

	serverCAPEM, err := os.ReadFile(cfg.ServerCACert)
	if err != nil {
		return nil, fmt.Errorf("reading server CA %s: %w", cfg.ServerCACert, err)
	}

	return &KubeconfigGenerator{
		cfg:      cfg,
		caCert:   caCert,
		caKey:    caKey,
		serverCA: base64.StdEncoding.EncodeToString(serverCAPEM),
	}, nil
}

// Credential holds the raw PEM blobs and expiry for a generated client certificate.
// Used both for kubeconfig generation and the /credential exec-plugin endpoint.
type Credential struct {
	CertPEM []byte
	KeyPEM  []byte
	Expiry  time.Time
}

// GenerateCredential signs a fresh client certificate and returns the raw PEM data.
// Use this when you need the cert material directly (e.g. the exec credential plugin).
func (g *KubeconfigGenerator) GenerateCredential(username string, groups []string) (*Credential, error) {
	certPEM, keyPEM, err := g.signClientCert(username, groups)
	if err != nil {
		return nil, fmt.Errorf("signing cert for %s: %w", username, err)
	}
	block, _ := pem.Decode(certPEM)
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parsing signed cert: %w", err)
	}
	return &Credential{CertPEM: certPEM, KeyPEM: keyPEM, Expiry: cert.NotAfter}, nil
}

// Generate produces a kubeconfig with a freshly signed client certificate for username.
// groups are embedded as the certificate's Organisation field, which Kubernetes reads
// as RBAC group memberships.
func (g *KubeconfigGenerator) Generate(username string, groups []string) ([]byte, error) {
	cred, err := g.GenerateCredential(username, groups)
	if err != nil {
		return nil, err
	}
	return renderKubeconfig(
		g.cfg.ServerURL,
		g.serverCA,
		g.cfg.ClusterName,
		username,
		base64.StdEncoding.EncodeToString(cred.CertPEM),
		base64.StdEncoding.EncodeToString(cred.KeyPEM),
	)
}

// signClientCert issues an ECDSA P-256 client certificate signed by the k3s client CA.
func (g *KubeconfigGenerator) signClientCert(username string, groups []string) (certPEM, keyPEM []byte, err error) {
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("generating key: %w", err)
	}

	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, nil, fmt.Errorf("generating serial number: %w", err)
	}

	now := time.Now()
	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject: pkix.Name{
			CommonName:   username,
			Organization: groups, // Kubernetes maps these to RBAC groups
		},
		NotBefore:             now.Add(-5 * time.Minute), // tolerate minor clock skew
		NotAfter:              now.Add(g.cfg.Duration),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	certDER, err := x509.CreateCertificate(rand.Reader, tmpl, g.caCert, priv.Public(), g.caKey)
	if err != nil {
		return nil, nil, fmt.Errorf("signing certificate: %w", err)
	}
	certPEM = pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})

	privDER, err := x509.MarshalECPrivateKey(priv)
	if err != nil {
		return nil, nil, fmt.Errorf("marshaling private key: %w", err)
	}
	keyPEM = pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: privDER})

	return certPEM, keyPEM, nil
}

// loadCA reads a PEM certificate and its corresponding private key from disk.
// It handles EC, RSA, and PKCS#8 key formats.
func loadCA(certFile, keyFile string) (*x509.Certificate, crypto.PrivateKey, error) {
	certPEM, err := os.ReadFile(certFile)
	if err != nil {
		return nil, nil, fmt.Errorf("reading %s: %w", certFile, err)
	}
	block, _ := pem.Decode(certPEM)
	if block == nil {
		return nil, nil, fmt.Errorf("no PEM block in %s", certFile)
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, nil, fmt.Errorf("parsing certificate in %s: %w", certFile, err)
	}

	keyPEM, err := os.ReadFile(keyFile)
	if err != nil {
		return nil, nil, fmt.Errorf("reading %s: %w", keyFile, err)
	}
	keyBlock, _ := pem.Decode(keyPEM)
	if keyBlock == nil {
		return nil, nil, fmt.Errorf("no PEM block in %s", keyFile)
	}

	var key crypto.PrivateKey
	switch keyBlock.Type {
	case "EC PRIVATE KEY":
		key, err = x509.ParseECPrivateKey(keyBlock.Bytes)
	case "RSA PRIVATE KEY":
		key, err = x509.ParsePKCS1PrivateKey(keyBlock.Bytes)
	case "PRIVATE KEY":
		key, err = x509.ParsePKCS8PrivateKey(keyBlock.Bytes)
	default:
		return nil, nil, fmt.Errorf("unsupported key type %q in %s", keyBlock.Type, keyFile)
	}
	if err != nil {
		return nil, nil, fmt.Errorf("parsing key in %s: %w", keyFile, err)
	}
	return cert, key, nil
}

// GenerateBootstrap returns a kubeconfig that contains no user credentials —
// just the cluster endpoint, server CA, and an exec plugin stanza that will
// invoke "ward credential --server=<wardURL>" on demand.
// Distribute this file to users instead of manually constructing it.
func (g *KubeconfigGenerator) GenerateBootstrap(wardURL, username string) ([]byte, error) {
	var buf bytes.Buffer
	err := bootstrapTmpl.Execute(&buf, map[string]string{
		"Server":      g.cfg.ServerURL,
		"ServerCA":    g.serverCA,
		"WardURL": wardURL,
		"Username":    username,
		"Cluster":     g.cfg.ClusterName,
	})
	return buf.Bytes(), err
}

var bootstrapTmpl = template.Must(template.New("bootstrap").Parse(
	`apiVersion: v1
kind: Config
preferences: {}
clusters:
- cluster:
    certificate-authority-data: {{.ServerCA}}
    server: {{.Server}}
  name: {{.Cluster}}
contexts:
- context:
    cluster: {{.Cluster}}
    user: {{.Username}}
  name: {{.Cluster}}
current-context: {{.Cluster}}
users:
- name: {{.Username}}
  user:
    exec:
      apiVersion: client.authentication.k8s.io/v1
      command: ward
      args:
      - credential
      - --server={{.WardURL}}
      interactiveMode: IfAvailable
      provideClusterInfo: false
`))

var kubeconfigTmpl = template.Must(template.New("kubeconfig").Parse(
	`apiVersion: v1
kind: Config
preferences: {}
clusters:
- cluster:
    certificate-authority-data: {{.ServerCA}}
    server: {{.Server}}
  name: {{.Cluster}}
contexts:
- context:
    cluster: {{.Cluster}}
    user: {{.Username}}
  name: {{.Cluster}}
current-context: {{.Cluster}}
users:
- name: {{.Username}}
  user:
    client-certificate-data: {{.ClientCert}}
    client-key-data: {{.ClientKey}}
`))

func renderKubeconfig(server, serverCA, cluster, username, clientCert, clientKey string) ([]byte, error) {
	var buf bytes.Buffer
	err := kubeconfigTmpl.Execute(&buf, map[string]string{
		"Server":     server,
		"ServerCA":   serverCA,
		"Cluster":    cluster,
		"Username":   username,
		"ClientCert": clientCert,
		"ClientKey":  clientKey,
	})
	return buf.Bytes(), err
}