package rotator

import (
	"context"
	"crypto/ecdsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"net"
	"net/http"
	"strings"
	"time"
)

type vaultClient struct {
	baseURL string
	client  *http.Client
}

type vaultLoginResponse struct {
	Auth struct {
		ClientToken string `json:"client_token"`
	} `json:"auth"`
}

type vaultSignResponse struct {
	Data struct {
		Certificate string   `json:"certificate"`
		CAChain     []string `json:"ca_chain"`
	} `json:"data"`
}

type signedCertificate struct {
	leaf  *x509.Certificate
	chain []byte
	key   []byte
}

func (v vaultClient) login(ctx context.Context, role, jwt string) (string, error) {
	var response vaultLoginResponse
	_, err := jsonRequest(ctx, v.client, http.MethodPost, v.baseURL+"/v1/auth/kubernetes/login", map[string]string{
		"role": role,
		"jwt":  jwt,
	}, nil, &response)
	if err != nil {
		return "", fmt.Errorf("vault login: %w", err)
	}
	if response.Auth.ClientToken == "" {
		return "", errors.New("vault login returned no token")
	}
	return response.Auth.ClientToken, nil
}

func (v vaultClient) sign(ctx context.Context, path, token, trustedRootPath string, addresses []net.IP, ttl time.Duration, privateKey *ecdsa.PrivateKey, csr []byte) (signedCertificate, error) {
	ipSANs := make([]string, len(addresses))
	for index, address := range addresses {
		ipSANs[index] = address.String()
	}
	var response vaultSignResponse
	_, err := jsonRequest(ctx, v.client, http.MethodPost, v.baseURL+"/v1/"+strings.TrimLeft(path, "/"), map[string]any{
		"csr":     string(csr),
		"ip_sans": ipSANs,
		"ttl":     fmt.Sprintf("%ds", int64(ttl/time.Second)),
		"format":  "pem",
	}, map[string]string{"X-Vault-Token": token}, &response)
	if err != nil {
		return signedCertificate{}, fmt.Errorf("sign certificate: %w", err)
	}
	if response.Data.Certificate == "" || len(response.Data.CAChain) == 0 {
		return signedCertificate{}, errors.New("vault returned an incomplete certificate chain")
	}

	block, _ := pem.Decode([]byte(response.Data.Certificate))
	if block == nil || block.Type != "CERTIFICATE" {
		return signedCertificate{}, errors.New("vault returned an invalid leaf certificate")
	}
	leaf, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return signedCertificate{}, fmt.Errorf("parse leaf certificate: %w", err)
	}
	trustedRoots, err := certificatePool([]string{trustedRootPath})
	if err != nil {
		return signedCertificate{}, err
	}
	if err := verifySignedCertificate(leaf, response.Data.CAChain, trustedRoots, addresses, privateKey); err != nil {
		return signedCertificate{}, err
	}
	key, err := marshalPrivateKey(privateKey)
	if err != nil {
		return signedCertificate{}, err
	}

	chain := []byte(response.Data.Certificate)
	for _, certificate := range response.Data.CAChain[:len(response.Data.CAChain)-1] {
		if !strings.HasSuffix(string(chain), "\n") {
			chain = append(chain, '\n')
		}
		chain = append(chain, certificate...)
	}
	return signedCertificate{leaf: leaf, chain: chain, key: key}, nil
}

func verifySignedCertificate(leaf *x509.Certificate, chain []string, roots *x509.CertPool, addresses []net.IP, privateKey *ecdsa.PrivateKey) error {
	intermediates := x509.NewCertPool()
	for _, encoded := range chain {
		block, _ := pem.Decode([]byte(encoded))
		if block == nil || block.Type != "CERTIFICATE" {
			return errors.New("vault returned an invalid CA chain")
		}
		certificate, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			return fmt.Errorf("parse CA certificate: %w", err)
		}
		intermediates.AddCert(certificate)
	}
	if _, err := leaf.Verify(x509.VerifyOptions{
		Roots:         roots,
		Intermediates: intermediates,
		KeyUsages:     []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		DNSName:       addresses[0].String(),
	}); err != nil {
		return fmt.Errorf("verify signed certificate: %w", err)
	}
	if !hasExactIPSANs(leaf, addresses) || len(leaf.DNSNames) != 0 {
		return errors.New("signed certificate does not contain the exact required IP SAN set")
	}
	publicKey, ok := leaf.PublicKey.(*ecdsa.PublicKey)
	if !ok || !publicKey.Equal(&privateKey.PublicKey) {
		return errors.New("signed certificate does not match the generated key")
	}
	return nil
}

func hasExactIPSANs(certificate *x509.Certificate, expected []net.IP) bool {
	if len(certificate.IPAddresses) != len(expected) {
		return false
	}
	actual := make(map[string]struct{}, len(certificate.IPAddresses))
	for _, address := range certificate.IPAddresses {
		actual[address.String()] = struct{}{}
	}
	for _, address := range expected {
		if _, exists := actual[address.String()]; !exists {
			return false
		}
	}
	return true
}