package jwk

import (
	"crypto"
	"crypto/rsa"
	"encoding/binary"
	"fmt"
	"math/big"
	"reflect"

	"github.com/lestrrat-go/jwx/v3/internal/base64"
	"github.com/lestrrat-go/jwx/v3/internal/pool"
	"github.com/lestrrat-go/jwx/v3/jwa"
)

func init() {
	RegisterKeyExporter(jwa.RSA(), KeyExportFunc(rsaJWKToRaw))
}

func (k *rsaPrivateKey) Import(rawKey *rsa.PrivateKey) error {
	k.mu.Lock()
	defer k.mu.Unlock()

	d, err := bigIntToBytes(rawKey.D)
	if err != nil {
		return fmt.Errorf(`invalid rsa.PrivateKey: %w`, err)
	}
	k.d = d

	l := len(rawKey.Primes)

	if l < 0 /* I know, I'm being paranoid */ || l > 2 {
		return fmt.Errorf(`invalid number of primes in rsa.PrivateKey: need 0 to 2, but got %d`, len(rawKey.Primes))
	}

	if l > 0 {
		p, err := bigIntToBytes(rawKey.Primes[0])
		if err != nil {
			return fmt.Errorf(`invalid rsa.PrivateKey: %w`, err)
		}
		k.p = p
	}

	if l > 1 {
		q, err := bigIntToBytes(rawKey.Primes[1])
		if err != nil {
			return fmt.Errorf(`invalid rsa.PrivateKey: %w`, err)
		}
		k.q = q
	}

	// dp, dq, qi are optional values
	if v, err := bigIntToBytes(rawKey.Precomputed.Dp); err == nil {
		k.dp = v
	}
	if v, err := bigIntToBytes(rawKey.Precomputed.Dq); err == nil {
		k.dq = v
	}
	if v, err := bigIntToBytes(rawKey.Precomputed.Qinv); err == nil {
		k.qi = v
	}

	// public key part
	n, e, err := importRsaPublicKeyByteValues(&rawKey.PublicKey)
	if err != nil {
		return fmt.Errorf(`invalid rsa.PrivateKey: %w`, err)
	}
	k.n = n
	k.e = e

	return nil
}

func importRsaPublicKeyByteValues(rawKey *rsa.PublicKey) ([]byte, []byte, error) {
	n, err := bigIntToBytes(rawKey.N)
	if err != nil {
		return nil, nil, fmt.Errorf(`invalid rsa.PublicKey: %w`, err)
	}

	data := make([]byte, 8)
	binary.BigEndian.PutUint64(data, uint64(rawKey.E))
	i := 0
	for ; i < len(data); i++ {
		if data[i] != 0x0 {
			break
		}
	}
	return n, data[i:], nil
}

func (k *rsaPublicKey) Import(rawKey *rsa.PublicKey) error {
	k.mu.Lock()
	defer k.mu.Unlock()

	n, e, err := importRsaPublicKeyByteValues(rawKey)
	if err != nil {
		return fmt.Errorf(`invalid rsa.PrivateKey: %w`, err)
	}
	k.n = n
	k.e = e

	return nil
}

func buildRSAPublicKey(key *rsa.PublicKey, n, e []byte) {
	bin := pool.BigInt().Get()
	bie := pool.BigInt().Get()
	defer pool.BigInt().Put(bie)

	bin.SetBytes(n)
	bie.SetBytes(e)

	key.N = bin
	key.E = int(bie.Int64())
}

var rsaConvertibleKeys = []reflect.Type{
	reflect.TypeFor[RSAPrivateKey](),
	reflect.TypeFor[RSAPublicKey](),
}

func rsaJWKToRaw(key Key, hint any) (any, error) {
	extracted, err := extractEmbeddedKey(key, rsaConvertibleKeys)
	if err != nil {
		return nil, fmt.Errorf(`failed to extract embedded key: %w`, err)
	}
	switch key := extracted.(type) {
	case RSAPrivateKey:
		switch hint.(type) {
		case *rsa.PrivateKey, *any:
		default:
			return nil, fmt.Errorf(`invalid destination object type %T for private RSA JWK: %w`, hint, ContinueError())
		}

		locker, ok := key.(rlocker)
		if !ok {
			locker.rlock()
			defer locker.runlock()
		}

		od, ok := key.D()
		if !ok {
			return nil, fmt.Errorf(`missing "d" value`)
		}

		oq, ok := key.Q()
		if !ok {
			return nil, fmt.Errorf(`missing "q" value`)
		}

		op, ok := key.P()
		if !ok {
			return nil, fmt.Errorf(`missing "p" value`)
		}

		var d, q, p big.Int // note: do not use from sync.Pool

		d.SetBytes(od)
		q.SetBytes(oq)
		p.SetBytes(op)

		// optional fields
		var dp, dq, qi *big.Int

		if odp, ok := key.DP(); ok {
			dp = &big.Int{} // note: do not use from sync.Pool
			dp.SetBytes(odp)
		}

		if odq, ok := key.DQ(); ok {
			dq = &big.Int{} // note: do not use from sync.Pool
			dq.SetBytes(odq)
		}

		if oqi, ok := key.QI(); ok {
			qi = &big.Int{} // note: do not use from sync.Pool
			qi.SetBytes(oqi)
		}

		n, ok := key.N()
		if !ok {
			return nil, fmt.Errorf(`missing "n" value`)
		}

		e, ok := key.E()
		if !ok {
			return nil, fmt.Errorf(`missing "e" value`)
		}

		var privkey rsa.PrivateKey
		buildRSAPublicKey(&privkey.PublicKey, n, e)
		privkey.D = &d
		privkey.Primes = []*big.Int{&p, &q}

		if dp != nil {
			privkey.Precomputed.Dp = dp
		}
		if dq != nil {
			privkey.Precomputed.Dq = dq
		}
		if qi != nil {
			privkey.Precomputed.Qinv = qi
		}
		// This may look like a no-op, but it's required if we want to
		// compare it against a key generated by rsa.GenerateKey
		privkey.Precomputed.CRTValues = []rsa.CRTValue{}
		return &privkey, nil
	case RSAPublicKey:
		switch hint.(type) {
		case *rsa.PublicKey, *any:
		default:
			return nil, fmt.Errorf(`invalid destination object type %T for public RSA JWK: %w`, hint, ContinueError())
		}

		locker, ok := key.(rlocker)
		if !ok {
			locker.rlock()
			defer locker.runlock()
		}

		n, ok := key.N()
		if !ok {
			return nil, fmt.Errorf(`missing "n" value`)
		}

		e, ok := key.E()
		if !ok {
			return nil, fmt.Errorf(`missing "e" value`)
		}

		var pubkey rsa.PublicKey
		buildRSAPublicKey(&pubkey, n, e)

		return &pubkey, nil

	default:
		return nil, ContinueError()
	}
}

func makeRSAPublicKey(src Key) (Key, error) {
	newKey := newRSAPublicKey()

	// Iterate and copy everything except for the bits that should not be in the public key
	for _, k := range src.Keys() {
		switch k {
		case RSADKey, RSADPKey, RSADQKey, RSAPKey, RSAQKey, RSAQIKey:
			continue
		default:
			var v any
			if err := src.Get(k, &v); err != nil {
				return nil, fmt.Errorf(`rsa: makeRSAPublicKey: failed to get field %q: %w`, k, err)
			}
			if err := newKey.Set(k, v); err != nil {
				return nil, fmt.Errorf(`rsa: makeRSAPublicKey: failed to set field %q: %w`, k, err)
			}
		}
	}

	return newKey, nil
}

func (k *rsaPrivateKey) PublicKey() (Key, error) {
	return makeRSAPublicKey(k)
}

func (k *rsaPublicKey) PublicKey() (Key, error) {
	return makeRSAPublicKey(k)
}

// Thumbprint returns the JWK thumbprint using the indicated
// hashing algorithm, according to RFC 7638
func (k rsaPrivateKey) Thumbprint(hash crypto.Hash) ([]byte, error) {
	k.mu.RLock()
	defer k.mu.RUnlock()

	var key rsa.PrivateKey
	if err := Export(&k, &key); err != nil {
		return nil, fmt.Errorf(`failed to export RSA private key: %w`, err)
	}
	return rsaThumbprint(hash, &key.PublicKey)
}

func (k rsaPublicKey) Thumbprint(hash crypto.Hash) ([]byte, error) {
	k.mu.RLock()
	defer k.mu.RUnlock()

	var key rsa.PublicKey
	if err := Export(&k, &key); err != nil {
		return nil, fmt.Errorf(`failed to export RSA public key: %w`, err)
	}
	return rsaThumbprint(hash, &key)
}

func rsaThumbprint(hash crypto.Hash, key *rsa.PublicKey) ([]byte, error) {
	buf := pool.BytesBuffer().Get()
	defer pool.BytesBuffer().Put(buf)

	buf.WriteString(`{"e":"`)
	buf.WriteString(base64.EncodeUint64ToString(uint64(key.E)))
	buf.WriteString(`","kty":"RSA","n":"`)
	buf.WriteString(base64.EncodeToString(key.N.Bytes()))
	buf.WriteString(`"}`)

	h := hash.New()
	if _, err := buf.WriteTo(h); err != nil {
		return nil, fmt.Errorf(`failed to write rsaThumbprint: %w`, err)
	}
	return h.Sum(nil), nil
}

func validateRSAKey(key interface {
	N() ([]byte, bool)
	E() ([]byte, bool)
}, checkPrivate bool) error {
	n, ok := key.N()
	if !ok {
		return fmt.Errorf(`missing "n" value`)
	}

	e, ok := key.E()
	if !ok {
		return fmt.Errorf(`missing "e" value`)
	}

	if len(n) == 0 {
		// Ideally we would like to check for the actual length, but unlike
		// EC keys, we have nothing in the key itself that will tell us
		// how many bits this key should have.
		return fmt.Errorf(`missing "n" value`)
	}
	if len(e) == 0 {
		return fmt.Errorf(`missing "e" value`)
	}
	if checkPrivate {
		if priv, ok := key.(keyWithD); ok {
			if d, ok := priv.D(); !ok || len(d) == 0 {
				return fmt.Errorf(`missing "d" value`)
			}
		} else {
			return fmt.Errorf(`missing "d" value`)
		}
	}

	return nil
}

func (k *rsaPrivateKey) Validate() error {
	if err := validateRSAKey(k, true); err != nil {
		return NewKeyValidationError(fmt.Errorf(`jwk.RSAPrivateKey: %w`, err))
	}
	return nil
}

func (k *rsaPublicKey) Validate() error {
	if err := validateRSAKey(k, false); err != nil {
		return NewKeyValidationError(fmt.Errorf(`jwk.RSAPublicKey: %w`, err))
	}
	return nil
}