package signature

import (
	"bytes"
	"crypto"
	"crypto/sha256"
	"encoding/base64"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"code.justin.tv/sse/malachai/pkg/jwtvalidation"
	"code.justin.tv/sse/malachai/pkg/sigutils"

	"github.com/SermoDigital/jose"
	joseCrypto "github.com/SermoDigital/jose/crypto"
	"github.com/SermoDigital/jose/jws"
	"github.com/gofrs/uuid"
)

// SignatureHeader is the name of the header that contains the s2s auth signature
const SignatureHeader = "Twitch-Service-Auth-Req-Signature"

// CanonicalRequestHashAlg is the hash alg we're using to hash canonical request
const CanonicalRequestHashAlg = "SHA-256"

// Custom claim names
const (
	CanonicalRequestHashAlgClaim = "canonical_request_hash_alg"
	CanonicalRequestHeadersClaim = "canonical_request_headers"
	CanonicalRequestHashClaim    = "canonical_request_hash"
)

// PrivateKeyStorer is a getter for the signing key, must be pem encoded
type PrivateKeyStorer func() (key crypto.Signer, fingerprint string, err error)

// NoSigningKeyError is returned when an error is returned trying to retrieve
// a private key
type NoSigningKeyError error

// Signer signs http requests
type Signer struct {
	CallerID   string
	Method     joseCrypto.SigningMethod
	PrivateKey PrivateKeyStorer
}

// calls s.PrivateKey and wraps error with NoSigningKeyError
func (s *Signer) privateKey() (key crypto.Signer, fingerprint string, err error) {
	key, fingerprint, err = s.PrivateKey()
	if err != nil {
		err = NoSigningKeyError(err)
	}
	return
}

// SignRequest signs an http request
func (s *Signer) SignRequest(r *http.Request) (err error) {
	return s.SignRequestWithHashedBody(r, nil)
}

// SignRequestWithHashedBody signs an http request using the hashed body
func (s *Signer) SignRequestWithHashedBody(r *http.Request, hashedBody []byte) (err error) {
	header, value, err := s.GenerateSignature(r, hashedBody)
	if err != nil {
		return err
	}
	r.Header.Set(header, string(value))
	return
}

// GenerateSignature returns the signed header and header name from the hashed body
func (s *Signer) GenerateSignature(r *http.Request, hashedBody []byte) (key string, value string, err error) {
	privateKey, fingerprint, err := s.privateKey()
	if err != nil {
		return "", "", err
	}

	entity := &jwtvalidation.SigningEntity{
		Caller:      s.CallerID,
		Fingerprint: fingerprint,
	}

	signature, err := generateSignature(r, hashedBody, s.Method, privateKey, entity)
	if err != nil {
		return "", "", err
	}

	return SignatureHeader, string(signature), nil
}

func generateSignature(r *http.Request, hashedBody []byte, signingMethod joseCrypto.SigningMethod, key crypto.Signer, caller *jwtvalidation.SigningEntity) (signature []byte, err error) {
	hashedCanonicalRequest, err := generateCanonicalRequestHash(r, hashedBody, nil)
	if err != nil {
		return
	}

	canonicalRequestHeaders := make([]string, 0, len(r.Header))
	for k := range r.Header {
		if _, ok := blacklistedHeaders[strings.ToLower(k)]; !ok {
			canonicalRequestHeaders = append(canonicalRequestHeaders, k)
		}
	}

	c := jws.Claims{
		CanonicalRequestHashAlgClaim: CanonicalRequestHashAlg,
		CanonicalRequestHeadersClaim: canonicalRequestHeaders,
		CanonicalRequestHashClaim:    base64.StdEncoding.EncodeToString(hashedCanonicalRequest),
	}

	now := jose.Now()

	id, err := uuid.NewV4()
	if err != nil {
		return
	}
	c.SetJWTID(id.String())
	c.SetNotBefore(now.Add(-5 * time.Minute))
	c.SetIssuedAt(now)
	c.SetExpiration(now.Add(5 * time.Minute))

	t := jws.NewJWT(c, signingMethod)
	t.(jws.JWS).Protected().Set(jwtvalidation.KeyIDHeaderName, caller.EncodeToString())
	signature, err = t.Serialize(key)

	return
}

func getURIPath(u *url.URL) string {
	var uri string

	if len(u.Opaque) > 0 {
		uri = "/" + strings.Join(strings.Split(u.Opaque, "/")[3:], "/")
	} else {
		uri = u.EscapedPath()
	}

	if len(uri) == 0 {
		uri = "/"
	}

	return uri
}

func generateCanonicalRequestHash(r *http.Request, hashedBody []byte, headers []string) (hash []byte, err error) {
	hasher := sha256.New()

	_, err = hasher.Write([]byte(r.Method))
	if err != nil {
		return
	}

	_, err = hasher.Write([]byte("\n"))
	if err != nil {
		return
	}

	_, err = hasher.Write([]byte(url.PathEscape(getURIPath(r.URL))))
	if err != nil {
		return
	}

	_, err = hasher.Write([]byte("\n"))
	if err != nil {
		return
	}

	qp, err := generateCanonicalQueryParams(r.URL.Query())
	if err != nil {
		return
	}

	_, err = hasher.Write([]byte(qp))
	if err != nil {
		return
	}

	_, err = hasher.Write([]byte("\n"))
	if err != nil {
		return
	}

	ch := newCanonicalHeader(r.Header)
	if headers != nil {
		ch = ch.OnlyHeaders(headers)
	}
	cHeaders, err := ch.Marshal()
	if err != nil {
		return
	}

	_, err = hasher.Write(cHeaders)
	if err != nil {
		return
	}

	_, err = hasher.Write([]byte("\n"))
	if err != nil {
		return
	}

	if hashedBody == nil {
		hashedBody, err = sigutils.GenerateHashedBody(r.Body)
		if err != nil {
			return
		}
	}
	_, err = hasher.Write(hashedBody)
	if err != nil {
		return
	}

	hash = hasher.Sum(nil)
	return
}

func generateCanonicalQueryParams(urlValues url.Values) (bs []byte, err error) {
	b := bytes.NewBuffer(nil)

	values := newKVItems(urlValues)
	sort.Sort(values)

	for iH, h := range values {
		if iH != 0 {
			_, err = b.WriteString("%26")
			if err != nil {
				return
			}
		}
		_, err = b.WriteString(url.QueryEscape(strings.ToLower(h.Key)))
		if err != nil {
			return
		}
		_, err = b.WriteString("%3D")
		if err != nil {
			return
		}
		_, err = b.WriteString(url.QueryEscape(h.Value))
		if err != nil {
			return
		}
	}

	bs = b.Bytes()
	return
}
