package frkoauth

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"regexp"
	"strings"
)

const (
	KEY_VERSION = "k7"
	TOK_VERSION = "v7"
)

type PrivateKey struct {
	KeyId string
	Key   ed25519.PrivateKey
}

type PublicKey struct {
	KeyId string
	Key   ed25519.PublicKey
}

func (privateKey *PrivateKey) PublicKey() *PublicKey {
	publicKey := privateKey.Key.Public().(ed25519.PublicKey)

	return &PublicKey{
		KeyId: privateKey.KeyId,
		Key:   publicKey,
	}
}

func (publicKey *PublicKey) Save() string {
	return KEY_VERSION + ".pub." + publicKey.KeyId + "." + base64.RawURLEncoding.EncodeToString(publicKey.Key)
}

func (privateKey *PrivateKey) Save() string {
	return KEY_VERSION + ".sec." + privateKey.KeyId + "." + base64.RawURLEncoding.EncodeToString(privateKey.Key)
}

func (privateKey *PrivateKey) Sign(plainText []byte) string {
	toSign := TOK_VERSION + "." + privateKey.KeyId + "." + base64.RawURLEncoding.EncodeToString([]byte(plainText))
	sigBytes := ed25519.Sign(privateKey.Key, []byte(toSign))

	return toSign + "." + base64.RawURLEncoding.EncodeToString(sigBytes)
}

func VerifyKeyId(keyId string) error {
	if len(keyId) != 16 {
		return errors.New("invalid key id length")
	}
	r := regexp.MustCompile(`^[a-zA-Z0-9-_]+$`)
	if !r.Match([]byte(keyId)) {
		return errors.New("invalid key id syntax")
	}

	return nil
}

func (publicKey *PublicKey) Verify(encodedText string) ([]byte, error) {
	tokenParts := strings.Split(encodedText, ".")
	if len(tokenParts) != 4 {
		return nil, errors.New("invalid token syntax")
	}
	if tokenParts[0] != TOK_VERSION {
		return nil, errors.New("invalid token version")
	}
	keyId := tokenParts[1]
	err := VerifyKeyId(keyId)
	if err != nil {
		return nil, err
	}
	if keyId != publicKey.KeyId {
		return nil, errors.New("unexpected key id")
	}
	signedText := tokenParts[0] + "." + tokenParts[1] + "." + tokenParts[2]
	encSigVal := tokenParts[3]
	sigBytes, err := base64.RawURLEncoding.DecodeString(encSigVal)
	if err != nil {
		return nil, errors.New("invalid signature syntax")
	}
	if len(sigBytes) != ed25519.SignatureSize {
		return nil, errors.New("invalid signature size")
	}
	if !ed25519.Verify(publicKey.Key, []byte(signedText), sigBytes) {
		return nil, errors.New("invalid signature")
	}

	plainText, err := base64.RawURLEncoding.DecodeString(tokenParts[2])
	if err != nil {
		return nil, err
	}

	return plainText, nil
}

func LoadPublicKey(keyStr string) (*PublicKey, error) {
	keyParts := strings.Split(keyStr, ".")
	if len(keyParts) != 4 {
		return nil, errors.New("malformed key")
	}
	if keyParts[0] != KEY_VERSION {
		return nil, errors.New("invalid key version")
	}
	if keyParts[1] != "pub" {
		return nil, errors.New("not a public key")
	}
	keyId := keyParts[2]
	err := VerifyKeyId(keyId)
	if err != nil {
		return nil, err
	}
	publicKeyBytes, err := base64.RawURLEncoding.DecodeString(keyParts[3])
	if err != nil {
		return nil, err
	}
	if len(publicKeyBytes) != ed25519.PublicKeySize {
		return nil, errors.New("invalid public key length")
	}
	edPubKey := ed25519.PublicKey(publicKeyBytes)
	publicKey := &PublicKey{
		KeyId: keyId,
		Key:   edPubKey,
	}

	return publicKey, nil
}

func LoadPrivateKey(keyStr string) (*PrivateKey, error) {
	keyParts := strings.Split(keyStr, ".")
	if len(keyParts) != 4 {
		return nil, errors.New("malformed key")
	}
	if keyParts[0] != KEY_VERSION {
		return nil, errors.New("invalid key version")
	}
	if keyParts[1] != "sec" {
		return nil, errors.New("not a private key")
	}
	keyId := keyParts[2]
	err := VerifyKeyId(keyId)
	if err != nil {
		return nil, err
	}
	privateKeyBytes, err := base64.RawURLEncoding.DecodeString(keyParts[3])
	if err != nil {
		return nil, err
	}
	if len(privateKeyBytes) != ed25519.PrivateKeySize {
		return nil, errors.New("invalid public key length")
	}
	edPrivKey := ed25519.PrivateKey(privateKeyBytes)
	privateKey := &PrivateKey{
		KeyId: keyId,
		Key:   edPrivKey,
	}

	return privateKey, nil
}

func Generate() (*PrivateKey, error) {
	keyIdBytes := make([]byte, 12)
	_, err := rand.Read(keyIdBytes)
	if err != nil {
		return nil, err
	}
	keyId := base64.RawURLEncoding.EncodeToString(keyIdBytes)
	_, privateKeyBytes, err := ed25519.GenerateKey(nil)
	if err != nil {
		return nil, err
	}
	privateKey := &PrivateKey{
		KeyId: keyId,
		Key:   privateKeyBytes,
	}

	return privateKey, nil
}
