package jwk

import (
	"crypto/rsa"
	"encoding/base64"
	"errors"
	"math/big"

	"git.daebt.dev/golang/auth/algo"
	"git.daebt.dev/golang/auth/algo/rs"
)

type TokenRSA struct {
	*Header
	E string `json:"e"`
	N string `json:"n"`
}

func (t *TokenRSA) encodeE(v int) {
	var (
		buf = []byte{}
		skp = true
	)

	for i := 56; i > -1; i -= 8 {
		b := byte(v >> i)
		if skp && b == 0 {
			continue
		}
		skp = false
		buf = append(buf, b)
	}

	t.E = base64.RawURLEncoding.EncodeToString(buf)
}

func (t *TokenRSA) decodeE() (int, error) {
	if t.E == "" {
		return 0, errors.New("e is empty")
	}

	buf, err := base64.RawURLEncoding.DecodeString(t.E)
	if err != nil {
		return 0, err
	}

	var (
		num int
		l   = len(buf) - 1
	)
	for i, b := range buf {
		i = (l - i) * 8
		num = num | int(b)<<i
		i += 8
	}

	return num, nil
}

func (t *TokenRSA) Algo() (algo.Algorithm, error) {
	n, err := base64.RawURLEncoding.DecodeString(t.N)
	if err != nil {
		return nil, err
	}

	e, err := t.decodeE()
	if err != nil {
		return nil, err
	}

	o := rs.WithPub(&rsa.PublicKey{
		N: new(big.Int).SetBytes(n),
		E: e,
	})

	switch t.Header.Algorithm {
	case rs.AlgorithmRS256:
		return rs.NewRS256(o)
	case rs.AlgorithmRS384:
		return rs.NewRS384(o)
	case rs.AlgorithmRS512:
		return rs.NewRS512(o)
	}

	return nil, errors.New("undefined algorithm type")
}

func (t *TokenRSA) KeyId() []byte {
	if t.Header.KeyId != "" {
		if val, err := base64.RawURLEncoding.DecodeString(t.Header.KeyId); err == nil {
			return val
		}
	}

	return nil
}

func NewRSA(kid []byte, v *rs.Algo) Token {
	pub := v.PublicKey()

	t := &TokenRSA{
		Header: &Header{
			KeyId:     "",
			KeyType:   v.Key(),
			Use:       UseSignature,
			Algorithm: v.Algo(),
		},
		E: "",
		N: base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
	}
	if len(kid) > 0 {
		t.Header.KeyId = base64.RawURLEncoding.EncodeToString(kid)
	}

	t.encodeE(pub.E)
	return t
}