package libsignal

/*
#include <signal_protocol.h>
#include <key_helper.h>
#include <session_pre_key.h>
#include <ratchet.h>
#include <curve.h>
*/
import "C"

import (
	"runtime"
	"time"
	"unsafe"
)

// IdentityKeyPair is the serialized (protobuf) form of a Curve25519
// identity key pair, as produced by ratchet_identity_key_pair_serialize.
// This is the form to persist; use SplitIdentityKeyPair/DecodeIdentityPublicKey
// to recover the individual keys when needed.
type IdentityKeyPair struct {
	Record []byte
}

// GenerateIdentityKeyPair creates a new identity key pair. Do this once per
// bridged-chat identity, at first use, and persist the result.
func GenerateIdentityKeyPair(ctx *Context) (*IdentityKeyPair, error) {
	var kp *C.ratchet_identity_key_pair
	if code := C.signal_protocol_key_helper_generate_identity_key_pair(&kp, ctx.raw); code != C.SG_SUCCESS {
		return nil, newError("signal_protocol_key_helper_generate_identity_key_pair", int(code))
	}
	defer C.signal_type_unref((*C.signal_type_base)(unsafe.Pointer(kp)))

	var buf *C.signal_buffer
	if code := C.ratchet_identity_key_pair_serialize(&buf, kp); code != C.SG_SUCCESS {
		return nil, newError("ratchet_identity_key_pair_serialize", int(code))
	}
	defer freeBuffer(buf)

	return &IdentityKeyPair{Record: bufferToBytes(buf)}, nil
}

func deserializeIdentityKeyPair(ctx *Context, identity *IdentityKeyPair) (*C.ratchet_identity_key_pair, error) {
	var kp *C.ratchet_identity_key_pair
	data := identity.Record
	var ptr *C.uint8_t
	if len(data) > 0 {
		ptr = (*C.uint8_t)(unsafe.Pointer(&data[0]))
	}
	code := C.ratchet_identity_key_pair_deserialize(&kp, ptr, C.size_t(len(data)), ctx.raw)
	runtime.KeepAlive(data)
	if code != C.SG_SUCCESS {
		return nil, newError("ratchet_identity_key_pair_deserialize", int(code))
	}
	return kp, nil
}

// SplitIdentityKeyPair deserializes a combined identity key pair record (as
// produced by GenerateIdentityKeyPair) into the separate public/private key
// buffers the library's identity-key-store callback (IdentityStore.
// GetIdentityKeyPair) is expected to return.
func SplitIdentityKeyPair(ctx *Context, record []byte) (public, private []byte, err error) {
	kp, err := deserializeIdentityKeyPair(ctx, &IdentityKeyPair{Record: record})
	if err != nil {
		return nil, nil, err
	}
	defer C.signal_type_unref((*C.signal_type_base)(unsafe.Pointer(kp)))

	var pubBuf *C.signal_buffer
	if code := C.ec_public_key_serialize(&pubBuf, C.ratchet_identity_key_pair_get_public(kp)); code != C.SG_SUCCESS {
		return nil, nil, newError("ec_public_key_serialize", int(code))
	}
	defer freeBuffer(pubBuf)

	var privBuf *C.signal_buffer
	if code := C.ec_private_key_serialize(&privBuf, C.ratchet_identity_key_pair_get_private(kp)); code != C.SG_SUCCESS {
		return nil, nil, newError("ec_private_key_serialize", int(code))
	}
	defer freeBuffer(privBuf)

	return bufferToBytes(pubBuf), bufferToBytes(privBuf), nil
}

// DecodeIdentityPublicKey extracts the raw public key bytes from a combined
// identity key pair record, for publishing in an XEP-0384 bundle.
func DecodeIdentityPublicKey(ctx *Context, record []byte) ([]byte, error) {
	kp, err := deserializeIdentityKeyPair(ctx, &IdentityKeyPair{Record: record})
	if err != nil {
		return nil, err
	}
	defer C.signal_type_unref((*C.signal_type_base)(unsafe.Pointer(kp)))

	var pubBuf *C.signal_buffer
	if code := C.ec_public_key_serialize(&pubBuf, C.ratchet_identity_key_pair_get_public(kp)); code != C.SG_SUCCESS {
		return nil, newError("ec_public_key_serialize", int(code))
	}
	defer freeBuffer(pubBuf)
	return bufferToBytes(pubBuf), nil
}

// GenerateRegistrationID creates a new registration id - a random number
// each identity picks once, at "install" time.
func GenerateRegistrationID(ctx *Context) (uint32, error) {
	var id C.uint32_t
	if code := C.signal_protocol_key_helper_generate_registration_id(&id, 0, ctx.raw); code != C.SG_SUCCESS {
		return 0, newError("signal_protocol_key_helper_generate_registration_id", int(code))
	}
	return uint32(id), nil
}

// PreKey is a serialized one-time prekey record plus its numeric id.
type PreKey struct {
	ID     uint32
	Record []byte
}

// GeneratePreKeys creates count one-time prekeys with sequential ids
// starting at start. Store every one; each is consumed (and should then be
// deleted) the first time a remote peer's session is built from it.
func GeneratePreKeys(ctx *Context, start, count uint32) ([]PreKey, error) {
	var head *C.signal_protocol_key_helper_pre_key_list_node
	if code := C.signal_protocol_key_helper_generate_pre_keys(&head, C.uint(start), C.uint(count), ctx.raw); code != C.SG_SUCCESS {
		return nil, newError("signal_protocol_key_helper_generate_pre_keys", int(code))
	}
	defer C.signal_protocol_key_helper_key_list_free(head)

	var preKeys []PreKey
	for node := head; node != nil; node = C.signal_protocol_key_helper_key_list_next(node) {
		pk := C.signal_protocol_key_helper_key_list_element(node)
		var buf *C.signal_buffer
		if code := C.session_pre_key_serialize(&buf, pk); code != C.SG_SUCCESS {
			return nil, newError("session_pre_key_serialize", int(code))
		}
		preKeys = append(preKeys, PreKey{
			ID:     uint32(C.session_pre_key_get_id(pk)),
			Record: bufferToBytes(buf),
		})
		freeBuffer(buf)
	}
	return preKeys, nil
}

func deserializePreKey(ctx *Context, record []byte) (*C.session_pre_key, error) {
	var pk *C.session_pre_key
	var ptr *C.uint8_t
	if len(record) > 0 {
		ptr = (*C.uint8_t)(unsafe.Pointer(&record[0]))
	}
	code := C.session_pre_key_deserialize(&pk, ptr, C.size_t(len(record)), ctx.raw)
	runtime.KeepAlive(record)
	if code != C.SG_SUCCESS {
		return nil, newError("session_pre_key_deserialize", int(code))
	}
	return pk, nil
}

// PreKeyInfo is the decoded, ready-to-publish material for a one-time prekey.
type PreKeyInfo struct {
	ID        uint32
	PublicKey []byte
}

// DecodePreKey deserializes a stored one-time prekey record (as produced by
// GeneratePreKeys) and extracts its id and raw public key, for publishing in
// an XEP-0384 bundle.
func DecodePreKey(ctx *Context, record []byte) (*PreKeyInfo, error) {
	pk, err := deserializePreKey(ctx, record)
	if err != nil {
		return nil, err
	}
	defer C.signal_type_unref((*C.signal_type_base)(unsafe.Pointer(pk)))

	var pubBuf *C.signal_buffer
	if code := C.ec_public_key_serialize(&pubBuf, C.ec_key_pair_get_public(C.session_pre_key_get_key_pair(pk))); code != C.SG_SUCCESS {
		return nil, newError("ec_public_key_serialize", int(code))
	}
	defer freeBuffer(pubBuf)

	return &PreKeyInfo{
		ID:        uint32(C.session_pre_key_get_id(pk)),
		PublicKey: bufferToBytes(pubBuf),
	}, nil
}

// SignedPreKey is a serialized signed prekey record plus its numeric id.
type SignedPreKey struct {
	ID     uint32
	Record []byte
}

// GenerateSignedPreKey creates a new signed prekey, signed by identity (in
// its combined serialized form, as produced by GenerateIdentityKeyPair),
// with the given id.
func GenerateSignedPreKey(ctx *Context, identity *IdentityKeyPair, id uint32) (*SignedPreKey, error) {
	idKeyPair, err := deserializeIdentityKeyPair(ctx, identity)
	if err != nil {
		return nil, err
	}
	defer C.signal_type_unref((*C.signal_type_base)(unsafe.Pointer(idKeyPair)))

	var signedPreKey *C.session_signed_pre_key
	timestamp := C.uint64_t(time.Now().UnixMilli())
	if code := C.signal_protocol_key_helper_generate_signed_pre_key(&signedPreKey, idKeyPair, C.uint32_t(id), timestamp, ctx.raw); code != C.SG_SUCCESS {
		return nil, newError("signal_protocol_key_helper_generate_signed_pre_key", int(code))
	}
	defer C.signal_type_unref((*C.signal_type_base)(unsafe.Pointer(signedPreKey)))

	var buf *C.signal_buffer
	if code := C.session_signed_pre_key_serialize(&buf, signedPreKey); code != C.SG_SUCCESS {
		return nil, newError("session_signed_pre_key_serialize", int(code))
	}
	defer freeBuffer(buf)

	return &SignedPreKey{
		ID:     uint32(C.session_signed_pre_key_get_id(signedPreKey)),
		Record: bufferToBytes(buf),
	}, nil
}

func deserializeSignedPreKey(ctx *Context, record []byte) (*C.session_signed_pre_key, error) {
	var spk *C.session_signed_pre_key
	var ptr *C.uint8_t
	if len(record) > 0 {
		ptr = (*C.uint8_t)(unsafe.Pointer(&record[0]))
	}
	code := C.session_signed_pre_key_deserialize(&spk, ptr, C.size_t(len(record)), ctx.raw)
	runtime.KeepAlive(record)
	if code != C.SG_SUCCESS {
		return nil, newError("session_signed_pre_key_deserialize", int(code))
	}
	return spk, nil
}

// SignedPreKeyInfo is the decoded, ready-to-publish material for a signed
// prekey: its numeric id, raw public key, and (both) signature forms the
// library maintains - a legacy signature and an OMEMO signature (libomemo-c
// computes both regardless of which protocol version ends up using it).
//
// When publishing a bundle (or constructing a RemoteBundle from one),
// picking the wrong one of these two for the target's protocol version
// fails verification with SG_ERR_INVALID_KEY, not a version-mismatch
// error: session_builder_process_pre_key_bundle re-serializes the signed
// prekey's public key via ec_public_key_serialize (version < 4, the
// 33-byte DJB_TYPE-prefixed legacy form) or ec_public_key_serialize_omemo
// (version >= 4, the raw 32-byte Montgomery form) and checks it against
// whichever signature is supplied - so use Signature for
// ProtocolVersionLegacy and SignatureOMEMO for ProtocolVersionModern.
type SignedPreKeyInfo struct {
	ID             uint32
	PublicKey      []byte
	Signature      []byte
	SignatureOMEMO []byte
}

// DecodeSignedPreKey deserializes a stored signed prekey record (as produced
// by GenerateSignedPreKey) and extracts the fields needed to publish it in
// an XEP-0384 bundle.
func DecodeSignedPreKey(ctx *Context, record []byte) (*SignedPreKeyInfo, error) {
	spk, err := deserializeSignedPreKey(ctx, record)
	if err != nil {
		return nil, err
	}
	defer C.signal_type_unref((*C.signal_type_base)(unsafe.Pointer(spk)))

	var pubBuf *C.signal_buffer
	if code := C.ec_public_key_serialize(&pubBuf, C.ec_key_pair_get_public(C.session_signed_pre_key_get_key_pair(spk))); code != C.SG_SUCCESS {
		return nil, newError("ec_public_key_serialize", int(code))
	}
	defer freeBuffer(pubBuf)

	return &SignedPreKeyInfo{
		ID:             uint32(C.session_signed_pre_key_get_id(spk)),
		PublicKey:      bufferToBytes(pubBuf),
		Signature:      cBytesToGo(C.session_signed_pre_key_get_signature(spk), C.session_signed_pre_key_get_signature_len(spk)),
		SignatureOMEMO: signedPreKeyOmemoSignature(spk),
	}, nil
}