package mls

// MLS key packages (RFC 9420 §10).

// KeyPackage provides public information: supported version/cipher suite,
// public keys, and credentials.
type KeyPackage struct {
	version     protocolVersion
	cipherSuite CipherSuite
	initKey     hpkePublicKey
	leafNode    leafNode
	extensions  []extension
	signature   []byte
}

func (pkg *KeyPackage) unmarshal(r *Reader) error {
	*pkg = KeyPackage{}

	v, ok := r.readUint16()
	if !ok {
		return errUnexpectedEOF
	}
	pkg.version = protocolVersion(v)
	if pkg.version != protocolVersionMLS10 {
		return errInvalidVersion
	}

	v, ok = r.readUint16()
	if !ok {
		return errUnexpectedEOF
	}
	pkg.cipherSuite = CipherSuite(v)

	pkg.initKey, ok = r.readOpaqueVec()
	if !ok {
		return errUnexpectedEOF
	}

	if err := pkg.leafNode.unmarshal(r); err != nil {
		return err
	}

	exts, err := unmarshalExtensionVec(r)
	if err != nil {
		return err
	}
	pkg.extensions = exts

	pkg.signature, ok = r.readOpaqueVec()
	if !ok {
		return errUnexpectedEOF
	}
	return nil
}

func (pkg *KeyPackage) marshalTBS(w *Writer) {
	w.addUint16(uint16(pkg.version))
	w.addUint16(uint16(pkg.cipherSuite))
	w.writeOpaqueVec([]byte(pkg.initKey))
	pkg.leafNode.marshal(w)
	marshalExtensionVec(w, pkg.extensions)
}

func (pkg *KeyPackage) marshal(w *Writer) {
	pkg.marshalTBS(w)
	w.writeOpaqueVec(pkg.signature)
}

// UnmarshalKeyPackage reads a key package from an MLS message envelope.
func UnmarshalKeyPackage(raw []byte) (*KeyPackage, error) {
	var msg mlsMessage
	if err := unmarshalRaw(raw, &msg); err != nil {
		return nil, err
	}
	if msg.wireFormat != wireFormatMLSKeyPackage {
		return nil, errInvalidWireFormat
	}
	return msg.keyPackage, nil
}

// UnmarshalRawKeyPackage reads a key package from bare TLS bytes (no envelope).
// This is the format used by NIP-EE (kind 443 events).
func UnmarshalRawKeyPackage(raw []byte) (*KeyPackage, error) {
	r := newReader(raw)
	pkg := &KeyPackage{}
	if err := pkg.unmarshal(&r); err != nil {
		return nil, err
	}
	if !r.empty() {
		return nil, errExcessBytes
	}
	return pkg, nil
}

// Bytes encodes the key package in an MLSMessage envelope.
func (pkg *KeyPackage) Bytes() []byte {
	raw, err := marshalRaw(&mlsMessage{
		version:    protocolVersionMLS10,
		wireFormat: wireFormatMLSKeyPackage,
		keyPackage: pkg,
	})
	if err != nil {
		panic("mls: failed to marshal key package")
	}
	return raw
}

// RawBytes encodes the key package as bare TLS bytes (no envelope).
func (pkg *KeyPackage) RawBytes() []byte {
	var w Writer
	pkg.marshal(&w)
	raw, err := w.bytes()
	if err != nil {
		panic("mls: failed to marshal raw key package")
	}
	return raw
}

// KeyPackageRef is a hash uniquely identifying a key package.
type KeyPackageRef []byte

func (ref KeyPackageRef) equal(other KeyPackageRef) bool {
	return bytesEqual([]byte(ref), []byte(other))
}

// PrivateKeyPackage holds private key material.
type PrivateKeyPackage struct {
	InitKey       []byte
	EncryptionKey []byte
	SignatureKey  []byte
}

// KeyPairPackage holds both public and private information.
type KeyPairPackage struct {
	Public  KeyPackage
	Private PrivateKeyPackage
}

// Marshal serializes the KeyPairPackage for persistence.
// Format: opaque_vec(publicRawBytes) || opaque_vec(InitKey) || opaque_vec(EncryptionKey) || opaque_vec(SignatureKey)
func (kpp *KeyPairPackage) Marshal() ([]byte, error) {
	var w Writer
	w.writeOpaqueVec(kpp.Public.RawBytes())
	w.writeOpaqueVec(kpp.Private.InitKey)
	w.writeOpaqueVec(kpp.Private.EncryptionKey)
	w.writeOpaqueVec(kpp.Private.SignatureKey)
	return w.bytes()
}

// UnmarshalKeyPairPackage restores a KeyPairPackage from bytes produced by Marshal.
func UnmarshalKeyPairPackage(raw []byte) (*KeyPairPackage, error) {
	r := newReader(raw)

	pubBytes, ok := r.readOpaqueVec()
	if !ok {
		return nil, errUnexpectedEOF
	}
	pub, err := UnmarshalRawKeyPackage(pubBytes)
	if err != nil {
		return nil, err
	}

	initKey, ok := r.readOpaqueVec()
	if !ok {
		return nil, errUnexpectedEOF
	}
	encKey, ok := r.readOpaqueVec()
	if !ok {
		return nil, errUnexpectedEOF
	}
	sigKey, ok := r.readOpaqueVec()
	if !ok {
		return nil, errUnexpectedEOF
	}

	return &KeyPairPackage{
		Public: *pub,
		Private: PrivateKeyPackage{
			InitKey:       initKey,
			EncryptionKey: encKey,
			SignatureKey:  sigKey,
		},
	}, nil
}

// KeyPackageOptions configures key package generation.
type KeyPackageOptions struct {
	CapabilityExtensions []extensionType
	LeafExtensions       []extension
	KeyPackageExtensions []extension
}