package jwt

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
)

var (
	//ErrTooShort represents an error where an input JWT has less than three sections.
	ErrTooShort = errors.New("jwt: invalid section count")
	//ErrInvalidSignature represents an error where a JWT has an invalid signature.
	ErrInvalidSignature = errors.New("jwt: invalid signature")
	//ErrInvalidHeader represents an error where a JWT's header does not match.
	ErrInvalidHeader = errors.New("jwt: invalid header")
	//ErrInvalidB64 represents an error where padding canonicalization fails due to invalid base64
	ErrInvalidB64 = errors.New("jwt: invalid base64 value")
)

//A Fragment represents the three base64 parts of a JSON Web Token
//canonicalized to base64url with padding.
type Fragment struct {
	Parts [3][]byte
	// the unmodified first 2 fields for signature checks
	// since the base64 itself is signed, padding must be the
	// same during signature checks.
	Prefix [2][]byte
}

//Header returns the base64url encoded jwt header.
func (f Fragment) Header() []byte { return f.Parts[0] }

//Claims returns the base64url encoded jwt claims.
func (f Fragment) Claims() []byte { return f.Parts[1] }

//Signature returns the base64url encoded jwt signature.
func (f Fragment) Signature() []byte { return f.Parts[2] }

var bDot = []byte(".")

//Parse parses a json web token into the standard form used by this package
//where the base64 components have padding as in encoding/base64.
func Parse(jwt []byte) (f Fragment, err error) {
	if copy(f.Parts[:], bytes.SplitN(jwt, bDot, -1)) < 3 {
		err = ErrTooShort
		return
	}

	copy(f.Prefix[:], f.Parts[:2])

	for i := range f.Parts {
		if f.Parts[i], err = canonicalizeb64(f.Parts[i]); err != nil {
			return
		}
	}

	return
}

var bMaxEqual = []byte("==")

//https://tools.ietf.org/html/rfc7515#appendix-C
func canonicalizeb64(i []byte) (o []byte, err error) {
	switch n := len(i) % 4; n {
	case 0: // no padding needed
		o = i
	case 2: // 2 chars needed
		o = make([]byte, len(i)+2)
		copy(o, i)
		copy(o[len(i):], bMaxEqual[:2])
	case 3: // 1 chars needed
		o = make([]byte, len(i)+1)
		copy(o, i)
		copy(o[len(i):], bMaxEqual[:1])
	default: // invalid
		err = ErrInvalidB64
	}

	return
}

//DecodeAndValidate decodes the json and validates the signature of the JSON Web Token
//in `jwt`. The Algorithm MUST be known in advance of the token, as the header
//cannot be trusted before the signature is validated.
//
//This is a convenience function short for:
//	var f Fragment
//	if f, err = Parse(jwt); err != nil {
//		return
//	}
//
//	if err = f.Validate(a); err != nil {
//		return
//	}
//
//	if err = f.Decode(header, claims); err != nil {
//		return
//	}
//
func DecodeAndValidate(header, claims interface{}, a Algorithm, jwt []byte) (err error) {
	var f Fragment
	if f, err = Parse(jwt); err != nil {
		return
	}

	if err = f.Validate(a); err != nil {
		return
	}

	if err = f.Decode(header, claims); err != nil {
		return
	}

	return
}

//Validate validates a JWT's signature with specified Algorithm. The Algorithm MUST be known in advance
//of the token, as the header cannot be trusted.
func (f Fragment) Validate(a Algorithm) (err error) {
	var sigbts []byte

	if sigbts, err = base64.URLEncoding.DecodeString(string(f.Signature())); err != nil {
		err = Err{err, fmt.Sprintf("in decoding signature b64 %s", f.Signature())}
		return
	}
	return a.Validate(bytes.Join(f.Prefix[:2], bDot), sigbts)
}

//Decode decodes the JWT's header and claims into the values
//indicated by the passed `header` and `claims` pointers.
func (f Fragment) Decode(header, claims interface{}) (err error) {

	if err = json.NewDecoder(
		base64.NewDecoder(base64.URLEncoding, bytes.NewReader(f.Header())),
	).Decode(header); err != nil {
		err = Err{err, "in decoding header b64 JSON"}
		return
	}

	if err = json.NewDecoder(
		base64.NewDecoder(base64.URLEncoding, bytes.NewReader(f.Claims())),
	).Decode(claims); err != nil {
		err = Err{err, "in decoding claims b64 JSON"}
		return
	}

	return

}
