package jwk

import (
	"crypto"
	"fmt"
	"reflect"
	"slices"

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

func init() {
	RegisterKeyExporter(KeyKind(jwa.OctetSeq().String()), KeyExportFunc(octetSeqToRaw))
}

func (k *symmetricKey) Import(rawKey []byte) error {
	k.mu.Lock()
	defer k.mu.Unlock()

	if len(rawKey) == 0 {
		return fmt.Errorf(`non-empty []byte key required`)
	}

	k.octets = slices.Clone(rawKey)

	return nil
}

var symmetricConvertibleKeys = []reflect.Type{
	reflect.TypeFor[SymmetricKey](),
}

func octetSeqToRaw(key Key, hint any) (any, error) {
	extracted, err := extractEmbeddedKey(key, symmetricConvertibleKeys)
	if err != nil {
		return nil, fmt.Errorf(`failed to extract embedded key: %w`, err)
	}

	switch key := extracted.(type) {
	case SymmetricKey:
		switch hint.(type) {
		case *[]byte, *any:
		default:
			return nil, fmt.Errorf(`invalid destination object type %T for symmetric key: %w`, hint, ContinueError())
		}

		// rlocker is unexported with unexported methods, so only our
		// concrete types implement it. A successful assertion lets us
		// type-assert to the concrete struct and read fields directly
		// under a single batch lock. This avoids nested RLock (which
		// deadlocks when a writer is pending) while preserving an
		// atomic snapshot of all fields.
		var ooctets []byte
		if locker, ok := key.(rlocker); ok {
			locker.rlock()
			concrete := key.(*symmetricKey) //nolint:forcetypeassert // rlocker is unexported; only our concrete types implement it
			ooctets = concrete.octets
			locker.runlock()
		} else {
			// External implementation — use self-locking interface getters.
			var ok bool
			if ooctets, ok = key.Octets(); !ok {
				return nil, fmt.Errorf(`jwk.SymmetricKey: missing "k" field`)
			}
		}

		if ooctets == nil {
			return nil, fmt.Errorf(`jwk.SymmetricKey: missing "k" field`)
		}

		octets := make([]byte, len(ooctets))
		copy(octets, ooctets)
		return octets, nil
	default:
		return nil, ContinueError()
	}
}

// Thumbprint returns the JWK thumbprint using the indicated
// hashing algorithm, according to RFC 7638
func (k *symmetricKey) Thumbprint(hash crypto.Hash) ([]byte, error) {
	k.mu.RLock()
	defer k.mu.RUnlock()
	var octets []byte
	if err := Export(k, &octets); err != nil {
		return nil, fmt.Errorf(`failed to export symmetric key: %w`, err)
	}

	h := hash.New()
	fmt.Fprint(h, `{"k":"`)
	fmt.Fprint(h, base64.EncodeToString(octets))
	fmt.Fprint(h, `","kty":"oct"}`)
	return h.Sum(nil), nil
}

func (k *symmetricKey) PublicKey() (Key, error) {
	newKey := newSymmetricKey()

	for _, key := range k.Keys() {
		var v any
		if err := k.Get(key, &v); err != nil {
			return nil, fmt.Errorf(`failed to get field %q: %w`, key, err)
		}

		if err := newKey.Set(key, v); err != nil {
			return nil, fmt.Errorf(`failed to set field %q: %w`, key, err)
		}
	}
	return newKey, nil
}

func (k *symmetricKey) Validate() error {
	octets, ok := k.Octets()
	if !ok || len(octets) == 0 {
		return NewKeyValidationError(fmt.Errorf(`jwk.SymmetricKey: missing "k" field`))
	}
	return nil
}