package jwk

import (
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rsa"
	"errors"
	"fmt"
	"math/big"
	"reflect"
	"sync"

	"github.com/lestrrat-go/blackmagic"
	"github.com/lestrrat-go/jwx/v3/internal/ecutil"
)

// # Converting between Raw Keys and `jwk.Key`s
//
// A converter that converts from a raw key to a `jwk.Key` is called a KeyImporter.
// A converter that converts from a `jwk.Key` to a raw key is called a KeyExporter.

var keyImporters = make(map[reflect.Type]KeyImporter)

// KeyKind identifies a key for exporter dispatch. Built-in key types
// use the key type string (e.g. "RSA", "EC", "OKP", "oct"). Keys that
// implement KeyKinder can return a more specific identity
// (e.g. "OKP:Ed448") to select a curve-specific exporter.
type KeyKind string

// KeyKinder is implemented by keys that need exporter dispatch
// beyond just their key type. For example, OKP keys return
// "OKP:<curve>" so that curve-specific exporters can be registered
// by external modules.
//
// Keys that do not implement this interface are dispatched by
// KeyType().String() alone.
type KeyKinder interface {
	KeyKind() KeyKind
}

var keyExporters = make(map[KeyKind][]KeyExporter)

var muKeyImporters sync.RWMutex
var muKeyExporters sync.RWMutex

// RegisterKeyImporter registers a KeyImporter for the given raw key.
// When `jwk.Import()` is called, the library looks up the importer for
// the given raw key type (via `reflect`) and executes it.
//
// Importer dispatch is single-valued per Go type: there is exactly
// one importer registered per `reflect.TypeOf(from)`. Registering a
// second importer for the same raw-key type silently replaces the
// previous entry — including built-in importers for
// `*rsa.PrivateKey`, `*ecdsa.PrivateKey`, and so on. Callers that
// need to guard against accidental overwrites should keep track of
// registrations themselves and avoid double-registration at init()
// time.
//
// This deliberately differs from the stacking behavior of
// [RegisterKeyExporter] (keyed by [KeyKind] strings) and
// [RegisterKeyParser] (an untyped-JSON fallback chain); importer
// dispatch is a single-value map keyed by Go type, with no
// equivalent dimension to try next. v4 turns the overwrite into
// an error; v3 keeps the frozen silent-overwrite behavior for
// backward compatibility.
func RegisterKeyImporter(from any, conv KeyImporter) {
	muKeyImporters.Lock()
	defer muKeyImporters.Unlock()
	keyImporters[reflect.TypeOf(from)] = conv
}

// RegisterKeyExporter registers a KeyExporter for the given key identity.
// When `jwk.Export()` is called, the library first tries exporters registered
// for the key's specific identity (via [KeyKinder]), then falls back to
// exporters registered for the key type alone.
//
// For most key types, pass `KeyKind(kty.String())` (e.g. `KeyKind("RSA")`).
// For curve-specific exporters, use a compound identity like `KeyKind("OKP:Ed448")`.
func RegisterKeyExporter(ident KeyKind, conv KeyExporter) {
	muKeyExporters.Lock()
	defer muKeyExporters.Unlock()
	convs, ok := keyExporters[ident]
	if !ok {
		convs = []KeyExporter{conv}
	} else {
		convs = append([]KeyExporter{conv}, convs...)
	}
	keyExporters[ident] = convs
}

// KeyImporter is used to convert from a raw key to a `jwk.Key`. mneumonic: from the PoV of the `jwk.Key`,
// we're _importing_ a raw key.
type KeyImporter interface {
	// Import takes the raw key to be converted, and returns a `jwk.Key` or an error if the conversion fails.
	Import(any) (Key, error)
}

// KeyImportFunc is a convenience type to implement KeyImporter as a function.
type KeyImportFunc func(any) (Key, error)

func (f KeyImportFunc) Import(raw any) (Key, error) {
	return f(raw)
}

// KeyExporter is used to convert from a `jwk.Key` to a raw key. mneumonic: from the PoV of the `jwk.Key`,
// we're _exporting_ it to a raw key.
type KeyExporter interface {
	// Export takes the `jwk.Key` to be converted, and a hint (the raw key to be converted to).
	// The hint is the object that the user requested the result to be assigned to.
	// The method should return the converted raw key, or an error if the conversion fails.
	//
	// Third party modules MUST NOT modifiy the hint object.
	//
	// When the user calls `key.Export(dst)`, the `dst` object is a _pointer_ to the
	// object that the user wants the result to be assigned to, but the converter
	// receives the _value_ that this pointer points to, to make it easier to
	// detect the type of the result.
	//
	// Note that the second argument may be an `any` (which means that the
	// user has delegated the type detection to the converter).
	//
	// Export must NOT modify the hint object, and should return jwk.ContinueError
	// if the hint object is not compatible with the converter.
	Export(Key, any) (any, error)
}

// KeyExportFunc is a convenience type to implement KeyExporter as a function.
type KeyExportFunc func(Key, any) (any, error)

func (f KeyExportFunc) Export(key Key, hint any) (any, error) {
	return f(key, hint)
}

func init() {
	{
		f := KeyImportFunc(rsaPrivateKeyToJWK)
		k := rsa.PrivateKey{}
		RegisterKeyImporter(k, f)
		RegisterKeyImporter(&k, f)
	}
	{
		f := KeyImportFunc(rsaPublicKeyToJWK)
		k := rsa.PublicKey{}
		RegisterKeyImporter(k, f)
		RegisterKeyImporter(&k, f)
	}
	{
		f := KeyImportFunc(ecdsaPrivateKeyToJWK)
		k := ecdsa.PrivateKey{}
		RegisterKeyImporter(k, f)
		RegisterKeyImporter(&k, f)
	}
	{
		f := KeyImportFunc(ecdsaPublicKeyToJWK)
		k := ecdsa.PublicKey{}
		RegisterKeyImporter(k, f)
		RegisterKeyImporter(&k, f)
	}
	{
		f := KeyImportFunc(okpPrivateKeyToJWK)
		for _, k := range []any{ed25519.PrivateKey(nil)} {
			RegisterKeyImporter(k, f)
		}
	}
	{
		f := KeyImportFunc(ecdhPrivateKeyToJWK)
		for _, k := range []any{ecdh.PrivateKey{}, &ecdh.PrivateKey{}} {
			RegisterKeyImporter(k, f)
		}
	}
	{
		f := KeyImportFunc(okpPublicKeyToJWK)
		for _, k := range []any{ed25519.PublicKey(nil)} {
			RegisterKeyImporter(k, f)
		}
	}
	{
		f := KeyImportFunc(ecdhPublicKeyToJWK)
		for _, k := range []any{ecdh.PublicKey{}, &ecdh.PublicKey{}} {
			RegisterKeyImporter(k, f)
		}
	}
	RegisterKeyImporter([]byte(nil), KeyImportFunc(bytesToKey))
}

func ecdhPrivateKeyToJWK(src any) (Key, error) {
	var raw *ecdh.PrivateKey
	switch src := src.(type) {
	case *ecdh.PrivateKey:
		raw = src
	case ecdh.PrivateKey:
		raw = &src
	default:
		return nil, fmt.Errorf(`cannot convert key type '%T' to ECDH jwk.Key`, src)
	}

	switch raw.Curve() {
	case ecdh.X25519():
		return okpPrivateKeyToJWK(raw)
	case ecdh.P256():
		return ecdhPrivateKeyToECJWK(raw, elliptic.P256())
	case ecdh.P384():
		return ecdhPrivateKeyToECJWK(raw, elliptic.P384())
	case ecdh.P521():
		return ecdhPrivateKeyToECJWK(raw, elliptic.P521())
	default:
		return nil, fmt.Errorf(`unsupported curve %s`, raw.Curve())
	}
}

func ecdhPrivateKeyToECJWK(raw *ecdh.PrivateKey, crv elliptic.Curve) (Key, error) {
	pub := raw.PublicKey()
	rawpub := pub.Bytes()

	size := ecutil.CalculateKeySize(crv)
	var x, y, d big.Int
	x.SetBytes(rawpub[1 : 1+size])
	y.SetBytes(rawpub[1+size:])
	d.SetBytes(raw.Bytes())

	var ecdsaPriv ecdsa.PrivateKey
	ecdsaPriv.Curve = crv
	ecdsaPriv.D = &d
	ecdsaPriv.X = &x
	ecdsaPriv.Y = &y
	return ecdsaPrivateKeyToJWK(&ecdsaPriv)
}

func ecdhPublicKeyToJWK(src any) (Key, error) {
	var raw *ecdh.PublicKey
	switch src := src.(type) {
	case *ecdh.PublicKey:
		raw = src
	case ecdh.PublicKey:
		raw = &src
	default:
		return nil, fmt.Errorf(`cannot convert key type '%T' to ECDH jwk.Key`, src)
	}

	switch raw.Curve() {
	case ecdh.X25519():
		return okpPublicKeyToJWK(raw)
	case ecdh.P256():
		return ecdhPublicKeyToECJWK(raw, elliptic.P256())
	case ecdh.P384():
		return ecdhPublicKeyToECJWK(raw, elliptic.P384())
	case ecdh.P521():
		return ecdhPublicKeyToECJWK(raw, elliptic.P521())
	default:
		return nil, fmt.Errorf(`unsupported curve %s`, raw.Curve())
	}
}

func ecdhPublicKeyToECJWK(raw *ecdh.PublicKey, crv elliptic.Curve) (Key, error) {
	rawbytes := raw.Bytes()
	size := ecutil.CalculateKeySize(crv)
	var x, y big.Int

	x.SetBytes(rawbytes[1 : 1+size])
	y.SetBytes(rawbytes[1+size:])
	var ecdsaPriv ecdsa.PublicKey
	ecdsaPriv.Curve = crv
	ecdsaPriv.X = &x
	ecdsaPriv.Y = &y
	return ecdsaPublicKeyToJWK(&ecdsaPriv)
}

// These may seem a bit repetitive and redandunt, but the problem is that
// each key type has its own Import method -- for example, Import(*ecdsa.PrivateKey)
// vs Import(*rsa.PrivateKey), and therefore they can't just be bundled into
// a single function.
func rsaPrivateKeyToJWK(src any) (Key, error) {
	var raw *rsa.PrivateKey
	switch src := src.(type) {
	case *rsa.PrivateKey:
		raw = src
	case rsa.PrivateKey:
		raw = &src
	default:
		return nil, fmt.Errorf(`cannot convert key type '%T' to RSA jwk.Key`, src)
	}
	k := newRSAPrivateKey()
	if err := k.Import(raw); err != nil {
		return nil, fmt.Errorf(`failed to initialize %T from %T: %w`, k, raw, err)
	}
	return k, nil
}

func rsaPublicKeyToJWK(src any) (Key, error) {
	var raw *rsa.PublicKey
	switch src := src.(type) {
	case *rsa.PublicKey:
		raw = src
	case rsa.PublicKey:
		raw = &src
	default:
		return nil, fmt.Errorf(`cannot convert key type '%T' to RSA jwk.Key`, src)
	}
	k := newRSAPublicKey()
	if err := k.Import(raw); err != nil {
		return nil, fmt.Errorf(`failed to initialize %T from %T: %w`, k, raw, err)
	}
	return k, nil
}

func ecdsaPrivateKeyToJWK(src any) (Key, error) {
	var raw *ecdsa.PrivateKey
	switch src := src.(type) {
	case *ecdsa.PrivateKey:
		raw = src
	case ecdsa.PrivateKey:
		raw = &src
	default:
		return nil, fmt.Errorf(`cannot convert key type '%T' to ECDSA jwk.Key`, src)
	}
	k := newECDSAPrivateKey()
	if err := k.Import(raw); err != nil {
		return nil, fmt.Errorf(`failed to initialize %T from %T: %w`, k, raw, err)
	}
	return k, nil
}

func ecdsaPublicKeyToJWK(src any) (Key, error) {
	var raw *ecdsa.PublicKey
	switch src := src.(type) {
	case *ecdsa.PublicKey:
		raw = src
	case ecdsa.PublicKey:
		raw = &src
	default:
		return nil, fmt.Errorf(`cannot convert key type '%T' to ECDSA jwk.Key`, src)
	}
	k := newECDSAPublicKey()
	if err := k.Import(raw); err != nil {
		return nil, fmt.Errorf(`failed to initialize %T from %T: %w`, k, raw, err)
	}
	return k, nil
}

func okpPrivateKeyToJWK(src any) (Key, error) {
	var raw any
	switch src.(type) {
	case ed25519.PrivateKey, *ecdh.PrivateKey:
		raw = src
	case ecdh.PrivateKey:
		raw = &src
	default:
		return nil, fmt.Errorf(`cannot convert key type '%T' to OKP jwk.Key`, src)
	}
	k := newOKPPrivateKey()
	if err := k.Import(raw); err != nil {
		return nil, fmt.Errorf(`failed to initialize %T from %T: %w`, k, raw, err)
	}
	return k, nil
}

func okpPublicKeyToJWK(src any) (Key, error) {
	var raw any
	switch src.(type) {
	case ed25519.PublicKey, *ecdh.PublicKey:
		raw = src
	case ecdh.PublicKey:
		raw = &src
	default:
		return nil, fmt.Errorf(`jwk: convert raw to OKP jwk.Key: cannot convert key type '%T' to OKP jwk.Key`, src)
	}
	k := newOKPPublicKey()
	if err := k.Import(raw); err != nil {
		return nil, fmt.Errorf(`failed to initialize %T from %T: %w`, k, raw, err)
	}
	return k, nil
}

func bytesToKey(src any) (Key, error) {
	var raw []byte
	switch src := src.(type) {
	case []byte:
		raw = src
	default:
		return nil, fmt.Errorf(`cannot convert key type '%T' to symmetric jwk.Key`, src)
	}

	k := newSymmetricKey()
	if err := k.Import(raw); err != nil {
		return nil, fmt.Errorf(`failed to initialize %T from %T: %w`, k, raw, err)
	}
	return k, nil
}

// Export converts a `jwk.Key` to a Export key. The dst argument must be a pointer to the
// object that the user wants the result to be assigned to.
//
// Normally you would pass a pointer to the zero value of the raw key type
// such as &(*rsa.PrivateKey) or &(*ecdsa.PublicKey), which gets assigned
// the converted key.
//
// If you do not know the exact type of a jwk.Key before attempting
// to obtain the raw key, you can simply pass a pointer to an
// empty interface as the second argument
//
// If you already know the exact type, it is recommended that you
// pass a pointer to the zero value of the actual key type for efficiency.
//
// Be careful when/if you are using a third party key type that implements
// the `jwk.Key` interface, as the first argument. This function tries hard
// to Do The Right Thing, but it is not guaranteed to work in all cases,
// especially when the object implements the `jwk.Key` interface via
// embedding.
func Export(key Key, dst any) error {
	if uk, ok := key.(UnsupportedKey); ok {
		kid, _ := uk.KeyID()
		return fmt.Errorf(`jwk.Export: cannot export an unsupported key (kty=%q, kid=%q) that could not be parsed; an extension module may be required: %w`, uk.KeyType().String(), kid, uk.Reason())
	}

	// dst better be a pointer
	rv := reflect.ValueOf(dst)
	if rv.Kind() != reflect.Ptr {
		return fmt.Errorf(`jwk.Export: destination object must be a pointer`)
	}

	muKeyExporters.RLock()
	exporters := findExporters(key)
	muKeyExporters.RUnlock()

	if len(exporters) == 0 {
		return fmt.Errorf(`jwk.Export: no exporters registered for key type '%T'`, key)
	}
	for _, conv := range exporters {
		v, err := conv.Export(key, dst)
		if err != nil {
			if errors.Is(err, ContinueError()) {
				continue
			}
			return fmt.Errorf(`jwk.Export: failed to export jwk.Key to raw format: %w`, err)
		}
		if err := blackmagic.AssignIfCompatible(dst, v); err != nil {
			return fmt.Errorf(`jwk.Export: failed to assign key: %w`, err)
		}
		return nil
	}
	return fmt.Errorf(`jwk.Export: no suitable exporter found for key type '%T'`, key)
}

// findExporters returns exporters for the key, trying the specific
// KeyKind first, then falling back to the key type. Caller must
// hold muKeyExporters.RLock.
func findExporters(key Key) []KeyExporter {
	if ki, ok := key.(KeyKinder); ok {
		ident := ki.KeyKind()
		if exporters, ok := keyExporters[ident]; ok {
			return exporters
		}
	}
	return keyExporters[KeyKind(key.KeyType().String())]
}