package station

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/pem"
	"errors"
	"io/fs"
	"os"
	"path/filepath"
)

const (
	caFileName   = "ca.pem"
	keyFileName  = "key.pem"
	csrFileName  = "csr.pem"
	certFileName = "cert.pem"
)

var (
	oidBasicConstraints = asn1.ObjectIdentifier{2, 5, 29, 19}
	oidKeyUsage         = asn1.ObjectIdentifier{2, 5, 29, 15}
	oidExtKeyUsage      = asn1.ObjectIdentifier{2, 5, 29, 37}
	oidClientAuth       = asn1.ObjectIdentifier{1, 3, 6, 1, 5, 5, 7, 3, 2}
)

func ensureCertDir(path string) error {
	err := os.Mkdir(path, 0o700)
	if err != nil {
		if !errors.Is(err, fs.ErrExist) {
			return err
		}
		info, statErr := os.Stat(path)
		if statErr != nil {
			return statErr
		}
		if !info.IsDir() {
			return err
		}
		return nil
	}
	return os.Chmod(path, 0o700)
}

func caDigest(dir string) (sum [32]byte, ok bool, err error) {
	body, err := os.ReadFile(filepath.Join(dir, caFileName))
	if errors.Is(err, fs.ErrNotExist) {
		return sum, false, nil
	}
	if err != nil {
		return sum, false, err
	}
	return sha256.Sum256(body), true, nil
}

func verifyCABundle(caPEM string, sig []byte, certPEM string) error {
	block, _ := pem.Decode([]byte(certPEM))
	if block == nil {
		return errors.New("signing certificate")
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return err
	}
	pub, ok := cert.PublicKey.(*rsa.PublicKey)
	if !ok {
		return errors.New("signing certificate")
	}
	sum := sha256.Sum256([]byte(caPEM))
	return rsa.VerifyPKCS1v15(pub, crypto.SHA256, sum[:], sig)
}

func writeCABundle(dir, caPEM string) error {
	return writeMode(filepath.Join(dir, caFileName), []byte(caPEM), 0o644)
}

func writeClientCert(dir, certPEM string) error {
	return writeMode(filepath.Join(dir, certFileName), []byte(certPEM), 0o644)
}

func csrNewerThanCert(dir string) bool {
	csr, err1 := os.Stat(filepath.Join(dir, csrFileName))
	cert, err2 := os.Stat(filepath.Join(dir, certFileName))
	if err1 != nil || err2 != nil {
		return false
	}
	return csr.ModTime().After(cert.ModTime())
}

func readCSR(dir string) (string, error) {
	body, err := os.ReadFile(filepath.Join(dir, csrFileName))
	if err != nil {
		return "", err
	}
	return string(body), nil
}

// prepareIdentity checks a saved certificate or creates the key and signing
// request. fatal is a window message when the station must stop.
func prepareIdentity(dir, callsign string) (tls.Certificate, bool, string) {
	certPath := filepath.Join(dir, certFileName)
	keyPath := filepath.Join(dir, keyFileName)
	_, certErr := os.Stat(certPath)
	_, keyErr := os.Stat(keyPath)
	certExists := certErr == nil
	keyExists := keyErr == nil
	if certErr != nil && !errors.Is(certErr, fs.ErrNotExist) {
		return tls.Certificate{}, false, "Cannot create " + certPath
	}
	if keyErr != nil && !errors.Is(keyErr, fs.ErrNotExist) {
		return tls.Certificate{}, false, "Cannot create " + keyPath
	}
	if certExists && !keyExists {
		return tls.Certificate{}, false, "Certificate key is missing"
	}
	if certExists && keyExists {
		pair, msg := loadPair(certPath, keyPath, callsign)
		if msg != "" {
			return tls.Certificate{}, false, msg
		}
		return pair, true, ""
	}
	if msg := ensureKeyAndCSR(dir, callsign, !keyExists); msg != "" {
		return tls.Certificate{}, false, msg
	}
	return tls.Certificate{}, false, ""
}

func loadPair(certPath, keyPath, callsign string) (tls.Certificate, string) {
	cert, err := loadCert(certPath)
	if err != nil {
		return tls.Certificate{}, "Certificate does not match the key"
	}
	key, err := loadKey(keyPath)
	if err != nil {
		return tls.Certificate{}, "Certificate does not match the key"
	}
	pub, ok := cert.PublicKey.(*rsa.PublicKey)
	if !ok || !pub.Equal(&key.PublicKey) {
		return tls.Certificate{}, "Certificate does not match the key"
	}
	if cert.Subject.CommonName != callsign {
		return tls.Certificate{}, "Certificate callsign does not match"
	}
	return tls.Certificate{
		Certificate: [][]byte{cert.Raw},
		PrivateKey:  key,
		Leaf:        cert,
	}, ""
}

func ensureKeyAndCSR(dir, callsign string, created bool) string {
	keyPath := filepath.Join(dir, keyFileName)
	var key *rsa.PrivateKey
	var err error
	if created {
		key, err = rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return "Cannot create " + keyPath
		}
		der := x509.MarshalPKCS1PrivateKey(key)
		body := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: der})
		if err := writeNew(keyPath, body, 0o600); err != nil {
			return "Cannot create " + keyPath
		}
	} else {
		key, err = loadKey(keyPath)
		if err != nil {
			return "Cannot create " + keyPath
		}
	}
	csrPath := filepath.Join(dir, csrFileName)
	_, csrErr := os.Stat(csrPath)
	if created || errors.Is(csrErr, fs.ErrNotExist) {
		if err := writeCSR(csrPath, key, callsign); err != nil {
			return "Cannot create " + csrPath
		}
	}
	return ""
}

func writeCSR(path string, key *rsa.PrivateKey, callsign string) error {
	bc, err := asn1.Marshal(struct {
		CA bool
	}{CA: false})
	if err != nil {
		return err
	}
	// digitalSignature, keyEncipherment, keyAgreement: bits 0, 2, and 4.
	ku, err := asn1.Marshal(asn1.BitString{Bytes: []byte{0xA8}, BitLength: 5})
	if err != nil {
		return err
	}
	eku, err := asn1.Marshal([]asn1.ObjectIdentifier{oidClientAuth})
	if err != nil {
		return err
	}
	der, err := x509.CreateCertificateRequest(rand.Reader, &x509.CertificateRequest{
		Subject: pkix.Name{CommonName: callsign},
		ExtraExtensions: []pkix.Extension{
			{Id: oidBasicConstraints, Critical: true, Value: bc},
			{Id: oidKeyUsage, Critical: true, Value: ku},
			{Id: oidExtKeyUsage, Critical: false, Value: eku},
		},
	}, key)
	if err != nil {
		return err
	}
	body := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: der})
	return writeMode(path, body, 0o644)
}

func loadCert(path string) (*x509.Certificate, error) {
	body, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	block, _ := pem.Decode(body)
	if block == nil {
		return nil, errors.New("certificate")
	}
	return x509.ParseCertificate(block.Bytes)
}

func loadKey(path string) (*rsa.PrivateKey, error) {
	body, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	block, _ := pem.Decode(body)
	if block == nil {
		return nil, errors.New("key")
	}
	switch block.Type {
	case "RSA PRIVATE KEY":
		return x509.ParsePKCS1PrivateKey(block.Bytes)
	case "PRIVATE KEY":
		k, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		rk, ok := k.(*rsa.PrivateKey)
		if !ok {
			return nil, errors.New("key")
		}
		return rk, nil
	default:
		return nil, errors.New("key")
	}
}

func writeNew(path string, body []byte, mode os.FileMode) error {
	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, mode)
	if err != nil {
		return err
	}
	_, werr := f.Write(body)
	cerr := f.Close()
	if werr != nil {
		return werr
	}
	if cerr != nil {
		return cerr
	}
	return os.Chmod(path, mode)
}

func writeMode(path string, body []byte, mode os.FileMode) error {
	if err := os.WriteFile(path, body, mode); err != nil {
		return err
	}
	return os.Chmod(path, mode)
}