package es256

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"encoding/json"
	"errors"
)

// ErrInvalidSig is returned when a signature is invalid.
var ErrInvalidSig = errors.New("Invalid Signature")

// PrivateKey is a JSON marshallable ECDSAPrivateKey
type PrivateKey struct {
	*ecdsa.PrivateKey
	KeyID string
}

// MarshalJSON implements json.Marshaler
func (epk PrivateKey) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct {
		KeyType   string    `json:"kty"`
		Use       string    `json:"use"`
		Curve     string    `json:"crv"`
		KeyOps    stringSet `json:"key_ops"`
		Algorithm string    `json:"alg"`
		KeyID     string    `json:"kid"`

		X *bigInt `json:"x"`
		Y *bigInt `json:"y"`
		D *bigInt `json:"d"`
	}{
		X: &bigInt{epk.PrivateKey.PublicKey.X},
		Y: &bigInt{epk.PrivateKey.PublicKey.Y},
		D: &bigInt{epk.PrivateKey.D},

		KeyType:   "EC",
		Use:       "sig",
		Curve:     "P-256",
		KeyOps:    *keyOpsSet,
		Algorithm: "ES256",
		KeyID:     epk.KeyID,
	})
}

// UnmarshalJSON implements json.Unmarshaler
func (epk *PrivateKey) UnmarshalJSON(inputJSON []byte) error {
	var raw struct {
		KeyType   string    `json:"kty"`
		Use       string    `json:"use"`
		Curve     string    `json:"crv"`
		KeyOps    stringSet `json:"key_ops"`
		Algorithm string    `json:"alg"`
		KeyID     string    `json:"kid"`

		X *bigInt `json:"x"`
		Y *bigInt `json:"y"`
		D *bigInt `json:"d"`
	}
	err := json.Unmarshal(inputJSON, &raw)
	if err != nil {
		return err
	}

	if raw.KeyType != "EC" {
		return errors.New("Incorrect KeyType")
	}
	if raw.Use != "sig" {
		return errors.New("Incorrect Use")
	}
	if raw.Curve != "P-256" {
		return errors.New("Incorrect Curve")
	}
	if raw.Algorithm != "ES256" {
		return errors.New("Incorrect Algorithm")
	}

	if !keyOpsSet.IsSubset(&raw.KeyOps) {
		return errors.New("Incorrect KeyOps")
	}

	epk.PrivateKey = new(ecdsa.PrivateKey)
	epk.PrivateKey.D = raw.D.Int
	epk.PrivateKey.PublicKey.Curve = elliptic.P256()
	epk.PrivateKey.PublicKey.X = raw.X.Int
	epk.PrivateKey.PublicKey.Y = raw.Y.Int
	epk.KeyID = raw.KeyID
	return nil
}

// Sign implements jwt.Algorithm
func (epk *PrivateKey) Sign(value []byte) ([]byte, error) {
	hashedVal := sha256.New()
	if _, err := hashedVal.Write(value); err != nil {
		return nil, err
	}
	r, s, err := ecdsa.Sign(rand.Reader, epk.PrivateKey, hashedVal.Sum(nil))

	if err != nil {
		return nil, err
	}

	ecdsaSig := &signature{
		R: r,
		S: s,
	}

	return ecdsaSig.Encode()
}

// Validate implements jwt.Algorithm
func (epk *PrivateKey) Validate(value, signature []byte) error {
	pk := &PublicKey{PublicKey: &epk.PrivateKey.PublicKey, KeyID: epk.KeyID}
	return pk.Validate(value, signature)
}

// Size implements jwt.Algorithm
func (epk *PrivateKey) Size() int {
	return 64
}

// Name implements jwt.Algorithm
func (epk *PrivateKey) Name() string {
	return "ES256"
}
