package jwt

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"reflect"
	"strings"
)

// JOSE represent a JSON Object Signing and Encryption header as
// defined in RFC 7515.
type JOSE struct {
	Algorithm            string   `json:"alg"`
	Type                 string   `json:"typ,omitempty"`
	Content              string   `json:"cty,omitempty"`
	Critical             []string `json:"crit,omitempty"`
	JWKSetURL            string   `json:"jku,omitempty"`
	JSONWebKey           string   `json:"jwk,omitempty"`
	KeyID                string   `json:"kid,omitempty"`
	X509URL              string   `json:"x5u,omitempty"`
	X509CertificateChain string   `json:"x5c,omitempty"`
	X509ThumbprintSha1   string   `json:"x5t,omitempty"`
	X509ThumbprintSha256 string   `json:"x5t#S256,omitempty"`

	//Sets of adfditional headers, private or public
	AdditionalHeaders map[string]interface{} `json:"-"`
}

// EncodeJSON encodes a JOSE header in JSON format. Not using
// MarshallJSON tyo avoid loops
func (j *JOSE) EncodeJSON() ([]byte, error) {
	data, err := json.Marshal(j)
	if err != nil {
		return nil, err
	}

	if len(data) == 0 || data[0] != '{' || data[len(data)-1] != '}' {
		return nil, fmt.Errorf("jws: Invalid JSON encoding: %s", data)
	}

	if len(j.AdditionalHeaders) == 0 {
		return data, nil
	}

	moreHeader, err := json.Marshal(j.AdditionalHeaders)
	if err != nil {
		return nil, err
	}

	if len(moreHeader) == 0 || moreHeader[0] != '{' || moreHeader[len(moreHeader)-1] != '}' {
		return nil, fmt.Errorf("jws: Invalid JSON encoding: %s", data)
	}
	data[len(data)-1] = ','
	return append(data, moreHeader[1:]...), nil
}

// EncodeBase64 encodes a JOSE into JSON base64 string. will allocate
// buffer.
func (j *JOSE) EncodeBase64() ([]byte, error) {
	dec, err := j.EncodeJSON()
	if err != nil {
		return nil, err
	}

	enc := make([]byte, Base64EncodedBufferLen(len(dec)))
	base64.URLEncoding.Encode(enc, dec)
	return enc[:Base64EncodedStrippedLen(len(dec))], nil
}

// DecodeJOSE a JOSE header from a base64 text as defined in the RFC 7515
func DecodeJOSE(data []byte) (*JOSE, error) {
	bData := make([]byte, Base64DecodedLenFromStripped(len(data)))

	if err := Base64Decode(bData, data); err != nil {
		return nil, err
	}

	j := &JOSE{}
	if err := json.NewDecoder(bytes.NewBuffer(bData)).Decode(j); err != nil {
		return nil, err
	}

	return j, nil
}

// Validate validates a JOSE header data. Maybe it should disappear
func (j *JOSE) Validate() error {
	if len(j.Algorithm) == 0 {
		return fmt.Errorf("jwt: missing 'alg' header")
	}

	return nil
}

type fieldSetter func(*JOSE, string)

var fieldSetters = make(map[string]fieldSetter)

// UnmarshalJSON is here to satisfy interface json.Unmarshaller. We
// need to provide ou own unmarshaller for the additional header.
func (j *JOSE) UnmarshalJSON(b []byte) error {
	raw := make(map[string]interface{})

	if err := json.Unmarshal(b, &raw); err != nil {
		return err
	}

	if j.AdditionalHeaders == nil {
		j.AdditionalHeaders = make(map[string]interface{}, len(raw))
	}

	for key, data := range raw {
		if fSetter, ok := fieldSetters[key]; ok == true {
			strData, ok := data.(string)
			if ok == false {
				return &json.UnmarshalTypeError{
					Value: reflect.TypeOf(data).Kind().String(),
					Type:  reflect.TypeOf(string("")),
				}
			}

			fSetter(j, strData)
			continue
		}

		j.AdditionalHeaders[key] = data
	}
	return nil
}

func init() {
	tJOSE := reflect.TypeOf(JOSE{})
	for i := 0; i < tJOSE.NumField(); i++ {
		jField := tJOSE.Field(i)
		jTag := jField.Tag.Get("json")
		if len(jTag) == 0 || jTag == "-" {
			continue
		}

		jName := strings.Split(jTag, ",")
		fieldSetters[jName[0]] = func(j *JOSE, data string) {
			reflect.ValueOf(j).Elem().FieldByName(jField.Name).SetString(data)
		}
	}
}