package crypto

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/sha512"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"strconv"

	"filippo.io/edwards25519"
)

// Byte sizes of key material.
const (
	SeedSize       = 32
	PublicKeySize  = 32
	PrivateKeySize = SeedSize + PublicKeySize
	signatureSize  = 64
	sha512Size     = 64
)

type PrivateKey struct {
	seed   [SeedSize]byte
	pub    [PublicKeySize]byte
	s      edwards25519.Scalar
	prefix [sha512Size / 2]byte
}

func (priv *PrivateKey) Bytes() []byte {
	k := make([]byte, 0, PrivateKeySize)
	k = append(k, priv.seed[:]...)
	k = append(k, priv.pub[:]...)
	return k
}

func (priv *PrivateKey) HexString() string {
	return hex.EncodeToString(priv.Bytes())
}

func (priv *PrivateKey) Seed() []byte {
	seed := priv.seed
	return seed[:]
}

func (priv *PrivateKey) Public() crypto.PublicKey {
	p := &PublicKey{}
	newPublicKey(p, priv.pub[:])
	return p
}

func (priv *PrivateKey) PublicKey() []byte {
	pub := priv.pub
	return pub[:]
}

type PublicKey struct {
	a      edwards25519.Point
	aBytes [32]byte
}

func (pub *PublicKey) Bytes() []byte {
	a := pub.aBytes
	return a[:]
}

func (pub *PublicKey) String() string {
	return hex.EncodeToString(pub.Bytes())
}

func (pub *PublicKey) EqualTo(other *PublicKey) bool {
	if pub == nil || other == nil {
		return false
	}
	return bytes.Equal(pub.aBytes[:], other.aBytes[:])
}

func (pub *PublicKey) MarshalJSON() ([]byte, error) {
	return json.Marshal(hex.EncodeToString(pub.Bytes()))
}

func (pub *PublicKey) UnmarshalJSON(data []byte) error {
	var s string
	if err := json.Unmarshal(data, &s); err != nil {
		return err
	}
	b, err := hex.DecodeString(s)
	if err != nil {
		return err
	}
	k, err := NewPublicKey(b)
	if err != nil {
		return err
	}

	pub.a = k.a
	copy(pub.aBytes[:], k.aBytes[:])
	return nil
}

// GenerateKey generates a new Ed25519 private key pair.
func GenerateKey() (*PublicKey, *PrivateKey, error) {
	priv := &PrivateKey{}
	key, err := generateKey(priv)
	if err != nil {
		return nil, nil, err
	}
	pub, err := NewPublicKey(priv.PublicKey())
	return pub, key, err
}

func generateKey(priv *PrivateKey) (*PrivateKey, error) {
	rand.Read(priv.seed[:])
	precomputePrivateKey(priv)
	return priv, nil
}

func NewPrivateKeyFromSeed(seed []byte) (*PrivateKey, error) {
	priv := &PrivateKey{}
	return newPrivateKeyFromSeed(priv, seed)
}

func newPrivateKeyFromSeed(priv *PrivateKey, seed []byte) (*PrivateKey, error) {
	if l := len(seed); l != SeedSize {
		return nil, errors.New("ed25519: bad seed length: " + strconv.Itoa(l))
	}
	copy(priv.seed[:], seed)
	precomputePrivateKey(priv)
	return priv, nil
}

func DecodePrivateKey(s string) (*PrivateKey, error) {
	var (
		b   []byte
		err error
	)
	switch len(s) {
	case hex.EncodedLen(SeedSize):
		if b, err = hex.DecodeString(s); err != nil {
			return nil, err
		}
		return NewPrivateKeyFromSeed(b)
	case hex.EncodedLen(PrivateKeySize):
		if b, err = hex.DecodeString(s); err != nil {
			return nil, err
		}
		return NewPrivateKey(b)
	case base64.RawStdEncoding.EncodedLen(SeedSize):
		if b, err = base64.RawStdEncoding.DecodeString(s); err != nil {
			return nil, err
		}
		return NewPrivateKeyFromSeed(b)
	case base64.RawStdEncoding.EncodedLen(PrivateKeySize):
		if b, err = base64.RawStdEncoding.DecodeString(s); err != nil {
			return nil, err
		}
		return NewPrivateKey(b)
	case base64.StdEncoding.EncodedLen(SeedSize):
		if b, err = base64.StdEncoding.DecodeString(s); err != nil {
			return nil, err
		}
		return NewPrivateKeyFromSeed(b)
	case base64.StdEncoding.EncodedLen(PrivateKeySize):
		if b, err = base64.StdEncoding.DecodeString(s); err != nil {
			return nil, err
		}
		return NewPrivateKey(b)
	default:
		return nil, errors.New("ed25519: unknown private key encoding")
	}
}

func precomputePrivateKey(priv *PrivateKey) {
	hs := sha512.New()
	hs.Write(priv.seed[:])
	h := hs.Sum(make([]byte, 0, sha512Size))

	s, err := priv.s.SetBytesWithClamping(h[:32])
	if err != nil {
		panic("ed25519: internal error: setting scalar failed")
	}
	A := (&edwards25519.Point{}).ScalarBaseMult(s)
	copy(priv.pub[:], A.Bytes())

	copy(priv.prefix[:], h[32:])
}

func NewPrivateKey(priv []byte) (*PrivateKey, error) {
	p := &PrivateKey{}
	return newPrivateKey(p, priv)
}

func newPrivateKey(priv *PrivateKey, privBytes []byte) (*PrivateKey, error) {
	if l := len(privBytes); l != PrivateKeySize {
		return nil, errors.New("ed25519: bad private key length: " + strconv.Itoa(l))
	}

	copy(priv.seed[:], privBytes[:32])

	hs := sha512.New()
	hs.Write(priv.seed[:])
	h := hs.Sum(make([]byte, 0, sha512Size))

	if _, err := priv.s.SetBytesWithClamping(h[:32]); err != nil {
		panic("ed25519: internal error: setting scalar failed")
	}
	// Note that we are not decompressing the public key point here,
	// because it takes > 20% of the time of a signature generation.
	// Signing doesn't use it as a point anyway.
	copy(priv.pub[:], privBytes[32:])

	copy(priv.prefix[:], h[32:])

	return priv, nil
}

func DecodePublicKey(s string) (*PublicKey, error) {
	var (
		pub []byte
		err error
	)
	switch len(s) {
	case hex.EncodedLen(32):
		pub, err = hex.DecodeString(s)
	case base64.RawStdEncoding.EncodedLen(32):
		pub, err = base64.RawStdEncoding.DecodeString(s)
	case base64.StdEncoding.DecodedLen(32):
		pub, err = base64.StdEncoding.DecodeString(s)
	default:
		return nil, errors.New("auth: invalid public key")
	}
	if err != nil {
		return nil, err
	}

	return NewPublicKey(pub)
}

func NewPublicKey(pub []byte) (*PublicKey, error) {
	p := &PublicKey{}
	return newPublicKey(p, pub)
}

func newPublicKey(pub *PublicKey, pubBytes []byte) (*PublicKey, error) {
	if l := len(pubBytes); l != PublicKeySize {
		return nil, errors.New("ed25519: bad public key length: " + strconv.Itoa(l))
	}
	// SetBytes checks that the point is on the curve.
	if _, err := pub.a.SetBytes(pubBytes); err != nil {
		return nil, errors.New("ed25519: bad public key")
	}
	copy(pub.aBytes[:], pubBytes)
	return pub, nil
}

// Domain separation prefixes used to disambiguate Ed25519/Ed25519ph/Ed25519ctx.
// See RFC 8032, Section 2 and Section 5.1.
const (
	// domPrefixPure is empty for pure Ed25519.
	domPrefixPure = ""
	// domPrefixPh is dom2(phflag=1) for Ed25519ph. It must be followed by the
	// uint8-length prefixed context.
	domPrefixPh = "SigEd25519 no Ed25519 collisions\x01"
	// domPrefixCtx is dom2(phflag=0) for Ed25519ctx. It must be followed by the
	// uint8-length prefixed context.
	domPrefixCtx = "SigEd25519 no Ed25519 collisions\x00"
)

func Sign(priv *PrivateKey, message []byte) []byte {
	// Outline the function body so that the returned signature can be
	// stack-allocated.
	signature := make([]byte, signatureSize)
	return sign(signature, priv, message)
}

func sign(signature []byte, priv *PrivateKey, message []byte) []byte {
	return signWithDom(signature, priv, message, domPrefixPure, "")
}

func SignPH(priv *PrivateKey, message []byte, context string) ([]byte, error) {
	// Outline the function body so that the returned signature can be
	// stack-allocated.
	signature := make([]byte, signatureSize)
	return signPH(signature, priv, message, context)
}

func signPH(signature []byte, priv *PrivateKey, message []byte, context string) ([]byte, error) {
	if l := len(message); l != sha512Size {
		return nil, errors.New("ed25519: bad Ed25519ph message hash length: " + strconv.Itoa(l))
	}
	if l := len(context); l > 255 {
		return nil, errors.New("ed25519: bad Ed25519ph context length: " + strconv.Itoa(l))
	}
	return signWithDom(signature, priv, message, domPrefixPh, context), nil
}

func SignCtx(priv *PrivateKey, message []byte, context string) ([]byte, error) {
	// Outline the function body so that the returned signature can be
	// stack-allocated.
	signature := make([]byte, signatureSize)
	return signCtx(signature, priv, message, context)
}

func signCtx(signature []byte, priv *PrivateKey, message []byte, context string) ([]byte, error) {
	// FIPS 186-5 specifies Ed25519 and Ed25519ph (with context), but not Ed25519ctx.
	// Note that per RFC 8032, Section 5.1, the context SHOULD NOT be empty.
	if l := len(context); l > 255 {
		return nil, errors.New("ed25519: bad Ed25519ctx context length: " + strconv.Itoa(l))
	}
	return signWithDom(signature, priv, message, domPrefixCtx, context), nil
}

func signWithDom(signature []byte, priv *PrivateKey, message []byte, domPrefix, context string) []byte {
	mh := sha512.New()
	if domPrefix != domPrefixPure {
		mh.Write([]byte(domPrefix))
		mh.Write([]byte{byte(len(context))})
		mh.Write([]byte(context))
	}
	mh.Write(priv.prefix[:])
	mh.Write(message)
	messageDigest := make([]byte, 0, sha512Size)
	messageDigest = mh.Sum(messageDigest)
	r, err := edwards25519.NewScalar().SetUniformBytes(messageDigest)
	if err != nil {
		panic("ed25519: internal error: setting scalar failed")
	}

	R := (&edwards25519.Point{}).ScalarBaseMult(r)

	kh := sha512.New()
	if domPrefix != domPrefixPure {
		kh.Write([]byte(domPrefix))
		kh.Write([]byte{byte(len(context))})
		kh.Write([]byte(context))
	}
	kh.Write(R.Bytes())
	kh.Write(priv.pub[:])
	kh.Write(message)
	hramDigest := make([]byte, 0, sha512Size)
	hramDigest = kh.Sum(hramDigest)
	k, err := edwards25519.NewScalar().SetUniformBytes(hramDigest)
	if err != nil {
		panic("ed25519: internal error: setting scalar failed")
	}

	S := edwards25519.NewScalar().MultiplyAdd(k, &priv.s, r)

	copy(signature[:32], R.Bytes())
	copy(signature[32:], S.Bytes())

	return signature
}

func Verify(pub *PublicKey, message, sig []byte) error {
	return verify(pub, message, sig)
}

func verify(pub *PublicKey, message, sig []byte) error {
	return verifyWithDom(pub, message, sig, domPrefixPure, "")
}

func VerifyPH(pub *PublicKey, message []byte, sig []byte, context string) error {
	if l := len(message); l != sha512Size {
		return errors.New("ed25519: bad Ed25519ph message hash length: " + strconv.Itoa(l))
	}
	if l := len(context); l > 255 {
		return errors.New("ed25519: bad Ed25519ph context length: " + strconv.Itoa(l))
	}
	return verifyWithDom(pub, message, sig, domPrefixPh, context)
}

func VerifyCtx(pub *PublicKey, message []byte, sig []byte, context string) error {
	if l := len(context); l > 255 {
		return errors.New("ed25519: bad Ed25519ctx context length: " + strconv.Itoa(l))
	}
	return verifyWithDom(pub, message, sig, domPrefixCtx, context)
}

func verifyWithDom(pub *PublicKey, message, sig []byte, domPrefix, context string) error {
	if l := len(sig); l != signatureSize {
		return errors.New("ed25519: bad signature length: " + strconv.Itoa(l))
	}

	if sig[63]&224 != 0 {
		return errors.New("ed25519: invalid signature")
	}

	kh := sha512.New()
	if domPrefix != domPrefixPure {
		kh.Write([]byte(domPrefix))
		kh.Write([]byte{byte(len(context))})
		kh.Write([]byte(context))
	}
	kh.Write(sig[:32])
	kh.Write(pub.aBytes[:])
	kh.Write(message)
	hramDigest := make([]byte, 0, sha512Size)
	hramDigest = kh.Sum(hramDigest)
	k, err := edwards25519.NewScalar().SetUniformBytes(hramDigest)
	if err != nil {
		panic("ed25519: internal error: setting scalar failed")
	}

	S, err := edwards25519.NewScalar().SetCanonicalBytes(sig[32:])
	if err != nil {
		return errors.New("ed25519: invalid signature")
	}

	// [S]B = R + [k]A --> [k](-A) + [S]B = R
	minusA := (&edwards25519.Point{}).Negate(&pub.a)
	R := (&edwards25519.Point{}).VarTimeDoubleScalarBaseMult(k, minusA, S)

	if !bytes.Equal(sig[:32], R.Bytes()) {
		return errors.New("ed25519: invalid signature")
	}
	return nil
}

var _ crypto.PrivateKey = (*PrivateKey)(nil)