package omemo

import (
	"errors"
	"fmt"

	"dev.narayana.im/narayana/telegabber/e2ee"
	"dev.narayana.im/narayana/telegabber/e2ee/omemo/libsignal"
)

func (b *Backend) Encrypt(from e2ee.PeerID, to []e2ee.PeerID, plaintext []byte) (e2ee.Envelope, error) {
	login, owner, ok := e2ee.SplitOwnedPeer(from)
	if !ok {
		return e2ee.Envelope{}, fmt.Errorf("omemo: Encrypt: %q is not an OwnedPeer", from)
	}

	b.mu.Lock()
	defer b.mu.Unlock()

	store := b.db.Store(b.libctx, login, owner)

	// Outer content encryption happens once per message, regardless of how
	// many recipient devices it ends up going to - see envelope.go.
	payload, err := EncryptOuter(b.version, plaintext, owner)
	if err != nil {
		return e2ee.Envelope{}, fmt.Errorf("omemo: EncryptOuter: %w", err)
	}

	storeCtx, err := libsignal.NewStoreContext(b.libctx, store)
	if err != nil {
		return e2ee.Envelope{}, fmt.Errorf("omemo: NewStoreContext: %w", err)
	}
	defer storeCtx.Close()

	env := &WireEnvelope{
		Version:   b.version,
		SenderSID: ownDeviceIDNum,
		IV:        payload.IV,
		Payload:   payload.Ciphertext,
	}

	for _, peer := range to {
		deviceIDs, err := store.GetSubDeviceSessions(string(peer))
		if err != nil {
			return e2ee.Envelope{}, fmt.Errorf("omemo: GetSubDeviceSessions(%s): %w", peer, err)
		}
		if len(deviceIDs) == 0 {
			return e2ee.Envelope{}, fmt.Errorf("omemo: Encrypt: no known devices/sessions for %s (IngestRemoteBundle was not called)", peer)
		}

		for _, deviceID := range deviceIDs {
			cipher, err := libsignal.NewSessionCipher(b.libctx, storeCtx, libsignal.Address{Name: string(peer), DeviceID: deviceID}, b.version.libsignalVersion())
			if err != nil {
				return e2ee.Envelope{}, fmt.Errorf("omemo: NewSessionCipher(%s:%d): %w", peer, deviceID, err)
			}
			ct, err := cipher.Encrypt(payload.ContentKey)
			cipher.Close()
			if err != nil {
				return e2ee.Envelope{}, fmt.Errorf("omemo: Encrypt(%s:%d): %w", peer, deviceID, err)
			}
			env.Keys = append(env.Keys, WireKey{
				RecipientJID: string(peer),
				DeviceID:     deviceID,
				IsPreKey:     ct.Type == libsignal.CiphertextPreKeyType,
				Ciphertext:   ct.Serialized,
			})
		}
	}

	raw, err := env.Encode()
	if err != nil {
		return e2ee.Envelope{}, err
	}
	return e2ee.Envelope{Backend: Name, Raw: raw}, nil
}

func (b *Backend) Decrypt(from, to e2ee.PeerID, env e2ee.Envelope) ([]byte, error) {
	if env.Backend != Name {
		return nil, fmt.Errorf("omemo: Decrypt: envelope backend %q, want %q", env.Backend, Name)
	}
	login, owner, ok := e2ee.SplitOwnedPeer(to)
	if !ok {
		return nil, fmt.Errorf("omemo: Decrypt: %q is not an OwnedPeer", to)
	}

	wireEnv, err := DecodeWireEnvelope(env.Raw)
	if err != nil {
		return nil, fmt.Errorf("omemo: DecodeWireEnvelope: %w", err)
	}

	var keyForUs *WireKey
	for i := range wireEnv.Keys {
		if wireEnv.Keys[i].DeviceID == ownDeviceIDNum {
			keyForUs = &wireEnv.Keys[i]
			break
		}
	}
	if keyForUs == nil {
		return nil, errors.New("omemo: Decrypt: no <key> addressed to our own device in this envelope")
	}

	b.mu.Lock()
	defer b.mu.Unlock()

	store := b.db.Store(b.libctx, login, owner)
	regID, err := store.GetLocalRegistrationID()
	if err != nil {
		return nil, fmt.Errorf("omemo: GetLocalRegistrationID: %w", err)
	}

	storeCtx, err := libsignal.NewStoreContext(b.libctx, store)
	if err != nil {
		return nil, fmt.Errorf("omemo: NewStoreContext: %w", err)
	}
	defer storeCtx.Close()

	// The cipher's configured version must match what the SENDER actually
	// used (wireEnv.Version, derived from which XML namespace their
	// stanza used), not this backend's own outgoing preference (b.version)
	// - a remote peer's message version isn't guaranteed to match ours,
	// and this is the one case telegabber's own M1 round-trip test never
	// exercised (both sides there always shared the same configured
	// version), so this is deliberately conservative rather than assumed.
	cipher, err := libsignal.NewSessionCipher(b.libctx, storeCtx, libsignal.Address{Name: string(from), DeviceID: wireEnv.SenderSID}, wireEnv.Version.libsignalVersion())
	if err != nil {
		return nil, fmt.Errorf("omemo: NewSessionCipher: %w", err)
	}
	defer cipher.Close()

	msgType := libsignal.CiphertextSignalType
	if keyForUs.IsPreKey {
		msgType = libsignal.CiphertextPreKeyType
	}
	// Omemo1 and Omemo2 both use the OMEMOKeyExchange/OMEMOAuthenticatedMessage
	// wire framing (libsignal.ProtocolVersionV4); only Omemo0 uses the
	// original WhisperTextProtocol framing.
	useOmemoFraming := wireEnv.Version != Omemo0

	contentKey, err := cipher.Decrypt(b.libctx, libsignal.CiphertextMessage{Type: msgType, Serialized: keyForUs.Ciphertext}, useOmemoFraming, regID)
	if err != nil {
		return nil, fmt.Errorf("omemo: SessionCipher.Decrypt: %w", err)
	}

	plaintext, err := DecryptOuter(wireEnv.Version, &EncryptedPayload{Ciphertext: wireEnv.Payload, IV: wireEnv.IV, ContentKey: contentKey})
	if err != nil {
		return nil, fmt.Errorf("omemo: DecryptOuter: %w", err)
	}
	return plaintext, nil
}