package jwt

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"errors"
	"math/big"
)

var _ Algorithm = _ecc{}

var ErrInvalidECPrivateKey = errors.New("jwt: key is not a valid ECC private key")
var ErrSignatureLength = errors.New("jwt: incorrect length of signature")
var ErrInvalidECSignature = errors.New("jwt: signature invalid")

type _ecc struct {
	PrivateKey *ecdsa.PrivateKey
	Hash       crypto.Hash
	KeySize    int
	CurveBits  int
	name       string
}

func (e _ecc) Size() int { return e.CurveBits / 4 }

func (e _ecc) Name() string { return e.name }

func (e _ecc) Sign(value []byte) (signature []byte, err error) {
	hasher := e.Hash.New()
	hasher.Write(value)

	// The signature generated by ecdsa is represented a pair of big.Ints (r, s)
	r, s, err := ecdsa.Sign(rand.Reader, e.PrivateKey, hasher.Sum(nil))
	if err != nil {
		return nil, err
	}

	// Both arrays must be the length of key in bytes and the output must be 2* the length of the key in bytes.
	curveBits := e.PrivateKey.Curve.Params().BitSize
	if e.CurveBits != curveBits {
		return nil, ErrInvalidECPrivateKey
	}

	// Serialize big.Int outputs into big-endian byte arrays and pad them with zeros on
	// the left to make sure they are the correct size.
	keyBytes := curveBits / 8
	if curveBits%8 > 0 {
		keyBytes += 1
	}
	rBytes := r.Bytes()
	rBytesPadded := make([]byte, keyBytes)
	copy(rBytesPadded[keyBytes-len(rBytes):], rBytes)
	sBytes := s.Bytes()
	sBytesPadded := make([]byte, keyBytes)
	copy(sBytesPadded[keyBytes-len(sBytes):], sBytes)

	return append(rBytesPadded, sBytesPadded...), nil
}

func (e _ecc) Validate(body, signature []byte) (err error) {
	if len(signature) != 2*e.KeySize {
		return ErrSignatureLength
	}

	r := big.NewInt(0).SetBytes(signature[:e.KeySize])
	s := big.NewInt(0).SetBytes(signature[e.KeySize:])

	hasher := e.Hash.New()
	hasher.Write([]byte(body))

	// Verify the signature
	if verified := ecdsa.Verify(&e.PrivateKey.PublicKey, hasher.Sum(nil), r, s); verified == true {
		return nil
	}
	return ErrInvalidECSignature
}

func ES256(key *ecdsa.PrivateKey) Algorithm {
	return &_ecc{
		PrivateKey: key,
		name:       "ES256",
		Hash:       crypto.SHA256,
		KeySize:    32,
		CurveBits:  256,
	}
}

func ES384(key *ecdsa.PrivateKey) Algorithm {
	return &_ecc{
		PrivateKey: key,
		name:       "ES384",
		Hash:       crypto.SHA384,
		KeySize:    48,
		CurveBits:  384,
	}
}

func ES512(key *ecdsa.PrivateKey) Algorithm {
	return &_ecc{
		PrivateKey: key,
		name:       "ES512",
		Hash:       crypto.SHA512,
		KeySize:    66,
		CurveBits:  521,
	}
}

var eccAlgs = map[string]func(*ecdsa.PrivateKey) Algorithm{
	"ES256": ES256,
	"ES384": ES384,
	"ES512": ES512,
}
