package client

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"time"
)

// ErrCertificateExpired means a cached certificate can no longer be served.
var ErrCertificateExpired = errors.New("certificate expired")

// EnsureIdentity loads (or creates) the tunnel's private key and CSR. The CSR
// is reused byte-for-byte when it matches hostname, because binding a different
// CSR makes the server issue a new certificate.
func EnsureIdentity(paths Paths, hostname string) (keyPEM, csrPEM []byte, err error) {
	if err := os.MkdirAll(paths.Dir, 0o700); err != nil {
		return nil, nil, err
	}
	keyPEM, err = os.ReadFile(paths.Key)
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return nil, nil, err
	}
	var key *ecdsa.PrivateKey
	if len(keyPEM) > 0 {
		key, err = parseKey(keyPEM)
		if err != nil {
			return nil, nil, fmt.Errorf("private key: %w", err)
		}
	}
	if key == nil {
		key, err = ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
		if err != nil {
			return nil, nil, err
		}
		der, err := x509.MarshalPKCS8PrivateKey(key)
		if err != nil {
			return nil, nil, err
		}
		keyPEM = pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der})
		if err := writeAtomic(paths.Key, keyPEM, 0o600); err != nil {
			return nil, nil, err
		}
	}

	csrPEM, err = os.ReadFile(paths.CSR)
	if err == nil && csrMatches(csrPEM, hostname) {
		return keyPEM, csrPEM, nil
	}
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return nil, nil, err
	}
	csrPEM, err = generateCSR(key, hostname)
	if err != nil {
		return nil, nil, err
	}
	if err := writeAtomic(paths.CSR, csrPEM, 0o600); err != nil {
		return nil, nil, err
	}
	return keyPEM, csrPEM, nil
}

func parseKey(keyPEM []byte) (*ecdsa.PrivateKey, error) {
	block, _ := pem.Decode(keyPEM)
	if block == nil {
		return nil, errors.New("not a PEM private key")
	}
	parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	key, ok := parsed.(*ecdsa.PrivateKey)
	if !ok {
		return nil, errors.New("not an ECDSA private key")
	}
	if key.Curve != elliptic.P256() {
		return nil, errors.New("not a P-256 private key")
	}
	return key, nil
}

func generateCSR(key *ecdsa.PrivateKey, hostname string) ([]byte, error) {
	request := &x509.CertificateRequest{
		Subject:  pkix.Name{CommonName: hostname},
		DNSNames: []string{hostname, "*." + hostname},
	}
	der, err := x509.CreateCertificateRequest(rand.Reader, request, key)
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: der}), nil
}

func csrMatches(csrPEM []byte, hostname string) bool {
	block, _ := pem.Decode(csrPEM)
	if block == nil || block.Type != "CERTIFICATE REQUEST" {
		return false
	}
	csr, err := x509.ParseCertificateRequest(block.Bytes)
	if err != nil || csr.CheckSignature() != nil {
		return false
	}
	if csr.Subject.CommonName != hostname {
		return false
	}
	have := map[string]bool{}
	for _, name := range csr.DNSNames {
		have[name] = true
	}
	return have[hostname] && have["*."+hostname]
}

// LoadCertificate loads cert.pem/chain.pem/private-key.pem and checks that the
// certificate still covers hostname and has not expired.
func LoadCertificate(paths Paths, hostname string) (tls.Certificate, time.Time, error) {
	keyPEM, err := os.ReadFile(paths.Key)
	if err != nil {
		return tls.Certificate{}, time.Time{}, err
	}
	certPEM, err := os.ReadFile(paths.Cert)
	if err != nil {
		return tls.Certificate{}, time.Time{}, err
	}
	chainPEM, err := os.ReadFile(paths.Chain)
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return tls.Certificate{}, time.Time{}, err
	}
	return parseCertificate(append(certPEM, chainPEM...), keyPEM, hostname)
}

// SaveCertificate validates and stores a certificate from the API.
func SaveCertificate(paths Paths, certPEM, chainPEM, keyPEM []byte, hostname string) (tls.Certificate, time.Time, error) {
	fullchain := append(append([]byte{}, certPEM...), chainPEM...)
	cert, expiry, err := parseCertificate(fullchain, keyPEM, hostname)
	if err != nil {
		return tls.Certificate{}, time.Time{}, err
	}
	if err := writeAtomic(paths.Cert, certPEM, 0o644); err != nil {
		return tls.Certificate{}, time.Time{}, err
	}
	if err := writeAtomic(paths.Chain, chainPEM, 0o644); err != nil {
		return tls.Certificate{}, time.Time{}, err
	}
	return cert, expiry, nil
}

func parseCertificate(fullchainPEM, keyPEM []byte, hostname string) (tls.Certificate, time.Time, error) {
	cert, err := tls.X509KeyPair(fullchainPEM, keyPEM)
	if err != nil {
		return tls.Certificate{}, time.Time{}, fmt.Errorf("certificate/key: %w", err)
	}
	leaf, err := x509.ParseCertificate(cert.Certificate[0])
	if err != nil {
		return tls.Certificate{}, time.Time{}, err
	}
	if err := leaf.VerifyHostname(hostname); err != nil {
		return tls.Certificate{}, time.Time{}, fmt.Errorf("certificate does not cover %s: %w", hostname, err)
	}
	if time.Now().After(leaf.NotAfter) {
		return tls.Certificate{}, time.Time{}, ErrCertificateExpired
	}
	return cert, leaf.NotAfter, nil
}

// CertificateExpiry parses the API's RFC 3339 expiry string.
func CertificateExpiry(value string) (time.Time, error) {
	if value == "" {
		return time.Time{}, nil
	}
	return time.Parse(time.RFC3339, value)
}

func writeAtomic(path string, data []byte, mode os.FileMode) error {
	tmp, err := os.CreateTemp(filepath.Dir(path), ".tunnel-*.tmp")
	if err != nil {
		return err
	}
	name := tmp.Name()
	defer os.Remove(name)
	if err := tmp.Chmod(mode); err != nil {
		tmp.Close()
		return err
	}
	if _, err := tmp.Write(data); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	return os.Rename(name, path)
}