package shared

import (
	"crypto"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/x509"
	"encoding/pem"
	"errors"

	"git.mstar.dev/mstar/linstrom/config"
)

const sanityCheckRawMessage = "test message for sanity checking keys"

func GenerateKeypair(useEd bool) (publicKey []byte, privateKey []byte, err error) {
	if useEd {
		publicKey, privateKey, err := ed25519.GenerateKey(nil)
		if err != nil {
			return nil, nil, err
		}
		publicKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
		return publicKeyBytes, privateKey, nil

	} else {
		privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return nil, nil, err
		}
		privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
		publicKeyBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
		if err != nil {
			return nil, nil, err
		}
		return publicKeyBytes, privateKeyBytes, nil
	}
}

func Sign(toSign string, keyBytes []byte, keyIsRsa bool) ([]byte, error) {
	if keyIsRsa {
		key, err := x509.ParsePKCS1PrivateKey(keyBytes)
		if err != nil {
			return nil, err
		}
		hasher := sha256.New()
		_, err = hasher.Write([]byte(toSign))
		if err != nil {
			return nil, err
		}
		// hash := sha256.Sum256([]byte(toSign))
		hash := hasher.Sum(nil)
		signed, err := rsa.SignPKCS1v15(nil, key, crypto.SHA256, hash)
		// signed, err := key.Sign(rand.Reader, hash[:], crypto.SHA256)
		return signed, err
	} else {
		key := ed25519.PrivateKey(keyBytes)
		hash := sha512.Sum512([]byte(toSign))
		signed, err := key.Sign(rand.Reader, hash[:], crypto.SHA512)
		return signed, err
	}
}

func KeyBytesToPem(bytes []byte) string {
	var t string
	if config.GlobalConfig.Experimental.UseEd25519Keys {
		t = "PUBLIC KEY"
	} else {
		// t = "RSA PUBLIC KEY"
		t = "PUBLIC KEY"
	}
	block := pem.Block{
		Type:    t,
		Headers: nil,
		Bytes:   bytes,
	}
	return string(pem.EncodeToMemory(&block))
}

func SanityCheckRawEdKeys(pub ed25519.PublicKey, priv ed25519.PrivateKey) error {
	hash := sha512.Sum512([]byte(sanityCheckRawMessage))
	signed, err := priv.Sign(rand.Reader, hash[:], crypto.SHA512)
	if err != nil {
		return err
	}
	return ed25519.VerifyWithOptions(pub, hash[:], signed, &ed25519.Options{
		Hash: crypto.SHA512,
	})
}

func SanityCheckRawByteEdKeys(pub, priv []byte) error {
	pubKey := ed25519.PublicKey(pub)
	privKey := ed25519.PrivateKey(priv)
	return SanityCheckRawEdKeys(pubKey, privKey)
}

func SanityCheckX509dEdKeys(pub, priv []byte) error {
	privKey := ed25519.PrivateKey(priv)
	rawPubKey, err := x509.ParsePKIXPublicKey(pub)
	if err != nil {
		return err
	}
	pubKey, ok := rawPubKey.(ed25519.PublicKey)
	if !ok {
		return errors.New("not an ed25519 key")
	}
	return SanityCheckRawEdKeys(pubKey, privKey)
}

func SanityCheckPemdEdKeys(pub, priv []byte) error {
	privBlock, _ := pem.Decode(priv)
	pubBlock, _ := pem.Decode(pub)
	return SanityCheckX509dEdKeys(pubBlock.Bytes, privBlock.Bytes)
}

func SanityCheckRawRsaKeys(pub *rsa.PublicKey, priv *rsa.PrivateKey) error {
	hash := sha256.Sum256([]byte(sanityCheckRawMessage))
	signed, err := priv.Sign(rand.Reader, hash[:], crypto.SHA256)
	if err != nil {
		return err
	}
	return rsa.VerifyPKCS1v15(pub, crypto.SHA256, hash[:], signed)
}

func SanityCheckX509dRsaKeys(pub, priv []byte) error {
	privKey, err := x509.ParsePKCS1PrivateKey(priv)
	if err != nil {
		return err
	}
	rawPubKey, err := x509.ParsePKIXPublicKey(pub)
	if err != nil {
		return err
	}
	pubKey, ok := rawPubKey.(*rsa.PublicKey)
	if !ok {
		return errors.New("not an rsa key")
	}
	return SanityCheckRawRsaKeys(pubKey, privKey)
}

func SanityCheckPemdRsaKeys(pub, priv []byte) error {
	privBlock, _ := pem.Decode(priv)
	pubBlock, _ := pem.Decode(pub)
	return SanityCheckX509dRsaKeys(pubBlock.Bytes, privBlock.Bytes)
}