package common

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"strings"
	"time"

	"github.com/dcsunny/gocrypt"
	"github.com/lestrrat-go/jwx/jwa"
	"github.com/lestrrat-go/jwx/jwe"
)

type JWE struct {
	publicKey  *rsa.PublicKey
	privateKey *rsa.PrivateKey
	expire     int64
}

var (
	DefaultJWE *JWE
)

const (
	jweHeader = "eyJhbGciOiJSU0ExXzUiLCJlbmMiOiJBMTI4Q0JDLUhTMjU2In0"
)

//GenerateRsaKey 返回值 第一个private key,第二个是public key
func GenerateRsaKey() ([]byte, []byte, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, err
	}
	pk, err := x509.MarshalPKCS8PrivateKey(privateKey)
	if err != nil {
		return nil, nil, err
	}
	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: pk,
	}
	_pk := pem.EncodeToMemory(block)
	pb, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	if err != nil {
		return nil, nil, err
	}
	publicBlock := &pem.Block{Type: "RSA Public Key", Bytes: pb}
	_pb := pem.EncodeToMemory(publicBlock)
	return _pk, _pb, nil
}

func InitJWE(privateKey, publicKey string, expire int64) error {
	var err error
	DefaultJWE, err = NewJWE(privateKey, publicKey, expire)
	return err
}

func NewJWE(privateKey string, publicKey string, expire int64) (*JWE, error) {
	privateKey = strings.Replace(privateKey, "-----BEGIN RSA PRIVATE KEY-----", "", 1)
	privateKey = strings.Replace(privateKey, "-----END RSA PRIVATE KEY-----", "", 1)
	privateKey = strings.Replace(privateKey, "\n", "", -1)
	privateKeyDecoded, err := gocrypt.DecodeString(privateKey, gocrypt.Base64)
	pk, err := gocrypt.ParsePrivateKey(privateKeyDecoded, gocrypt.PKCS8)
	if err != nil {
		return nil, err
	}
	publicKey = strings.Replace(publicKey, "-----BEGIN RSA Public Key-----", "", 1)
	publicKey = strings.Replace(publicKey, "-----END RSA Public Key-----", "", 1)
	publicKey = strings.Replace(publicKey, "\n", "", -1)
	publicKeyDecoded, err := gocrypt.DecodeString(publicKey, gocrypt.Base64)
	pb, err := x509.ParsePKIXPublicKey(publicKeyDecoded)
	if err != nil {
		return nil, err
	}
	var j = new(JWE)
	j.publicKey = pb.(*rsa.PublicKey)
	j.privateKey = pk
	j.expire = expire
	return j, nil
}

type JWEClaims struct {
	AccountId string          `json:"id"`
	Expire    int64           `json:"exp"`
	Claims    json.RawMessage `json:"claims,omitempty"`
}

func (j *JWE) NewToken(accountID string, claims json.RawMessage) (string, error) {
	jweClaim := JWEClaims{
		AccountId: accountID,
		Claims:    claims,
		Expire:    time.Now().Unix() + j.expire,
	}

	payload, _ := json.Marshal(jweClaim)
	token, err := jwe.Encrypt(payload, jwa.RSA1_5, j.publicKey, jwa.A128CBC_HS256, jwa.NoCompress)
	if err != nil {
		return "", err
	}
	_token := string(token)
	_token = _token[52:]
	return _token, err
}

func (j *JWE) Parse(token string) (string, json.RawMessage, error) {
	token = jweHeader + "." + token
	decrypted, err := jwe.Decrypt([]byte(token), jwa.RSA1_5, j.privateKey)
	if err != nil {
		return "", nil, err
	}
	var jweClaim JWEClaims
	err = json.Unmarshal(decrypted, &jweClaim)
	if err != nil {
		return "", nil, err
	}
	if time.Now().Unix() > jweClaim.Expire {
		return "", nil, ErrTokenExpired
	}
	return jweClaim.AccountId, jweClaim.Claims, nil
}