package agentd

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"fmt"
	"io/fs"
	"os"

	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
	"google.golang.org/grpc"
)

// EnrollDialer opens a connection suitable for the (un-authenticated) Enroll RPC.
// The agent has no client certificate yet, so this connection presents none; the
// bootstrap token authenticates the call instead.
type EnrollDialer func(ctx context.Context) (*grpc.ClientConn, error)

// EnsureEnrolled guarantees the node holds a valid client certificate.
//
//   - If node.crt/node.key already exist, it loads and returns the certificate's
//     CN (the node UUID) without contacting the control plane (idempotent first
//     start vs. restart).
//   - Otherwise it generates an EC P-256 key + CSR, calls Enroll with the
//     bootstrap token, and persists node.key (0600), node.crt and ca.crt.
//
// The private key is generated on the node and never leaves it.
func EnsureEnrolled(ctx context.Context, cfg Config, dial EnrollDialer) (string, error) {
	cfg = cfg.withDefaults()

	if uuid, ok, err := loadEnrolledUUID(cfg); err != nil {
		return "", err
	} else if ok {
		return uuid, nil
	}

	if cfg.BootstrapToken == "" {
		return "", errors.New("agentd: not enrolled and no bootstrap token provided")
	}

	keyPEM, csrPEM, err := generateKeyAndCSR()
	if err != nil {
		return "", err
	}

	conn, err := dial(ctx)
	if err != nil {
		return "", fmt.Errorf("agentd: dial for enroll: %w", err)
	}
	defer conn.Close()

	client := agentv1.NewAgentServiceClient(conn)
	resp, err := client.Enroll(ctx, &agentv1.EnrollRequest{
		BootstrapToken: cfg.BootstrapToken,
		CSRPEM:         csrPEM,
		AgentVersion:   cfg.AgentVersion,
	})
	if err != nil {
		return "", fmt.Errorf("agentd: enroll rpc: %w", err)
	}
	if resp.NodeUUID == "" || len(resp.CertPEM) == 0 {
		return "", errors.New("agentd: enroll response missing node_uuid/cert")
	}

	if err := os.MkdirAll(cfg.StateDir, 0o700); err != nil {
		return "", fmt.Errorf("agentd: mkdir state dir: %w", err)
	}
	if err := atomicWrite(cfg.KeyPath(), keyPEM, 0o600); err != nil {
		return "", err
	}
	if err := atomicWrite(cfg.CertPath(), resp.CertPEM, 0o644); err != nil {
		return "", err
	}
	if len(resp.CAPEM) > 0 {
		if err := atomicWrite(cfg.CAPath(), resp.CAPEM, 0o644); err != nil {
			return "", err
		}
	}
	logf("enrolled successfully (cert valid until unix=%d)", resp.NotAfterUnix)
	return resp.NodeUUID, nil
}

// loadEnrolledUUID returns the CN of the persisted node certificate, if present.
func loadEnrolledUUID(cfg Config) (string, bool, error) {
	certPEM, err := os.ReadFile(cfg.CertPath())
	if errors.Is(err, fs.ErrNotExist) {
		return "", false, nil
	}
	if err != nil {
		return "", false, fmt.Errorf("agentd: read node cert: %w", err)
	}
	if _, err := os.Stat(cfg.KeyPath()); err != nil {
		// Cert without key is unusable — treat as not enrolled.
		return "", false, nil
	}
	block, _ := pem.Decode(certPEM)
	if block == nil {
		return "", false, errors.New("agentd: node cert has no PEM block")
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return "", false, fmt.Errorf("agentd: parse node cert: %w", err)
	}
	if cert.Subject.CommonName == "" {
		return "", false, errors.New("agentd: node cert has empty CN")
	}
	return cert.Subject.CommonName, true, nil
}

// generateKeyAndCSR creates an EC P-256 private key and a PKCS#10 CSR. The CSR's
// CN is a placeholder; the CA overrides it with the authoritative node UUID.
func generateKeyAndCSR() (keyPEM, csrPEM []byte, err error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("agentd: generate key: %w", err)
	}
	keyDER, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return nil, nil, fmt.Errorf("agentd: marshal key: %w", err)
	}
	keyPEM = pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})

	tmpl := &x509.CertificateRequest{
		Subject:            pkix.Name{CommonName: "pangolin-node-pending"},
		SignatureAlgorithm: x509.ECDSAWithSHA256,
	}
	csrDER, err := x509.CreateCertificateRequest(rand.Reader, tmpl, key)
	if err != nil {
		return nil, nil, fmt.Errorf("agentd: create CSR: %w", err)
	}
	csrPEM = pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: csrDER})
	return keyPEM, csrPEM, nil
}