package omemo

import (
	"crypto/rand"
	"errors"
	"fmt"
	"math/big"
	"strconv"
	"sync"

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

// Name identifies this backend in the e2ee registry and config.
const Name = "omemo"

// initialPreKeyCount is how many one-time prekeys EnsureIdentity generates
// up front - a conventional pool size (matches common OMEMO implementation
// practice) balancing bundle-refill frequency against bundle size.
// Replenishing a depleted pool (prekeys are removed automatically as
// they're consumed - see ListPreKeyIDs's doc comment) is a future
// enhancement, not handled yet.
const initialPreKeyCount = 100

// ownDeviceIDNum is the sole device id this gateway ever publishes for an
// owned identity - there is exactly one telegabber process encrypting on
// behalf of any given bridged chat pseudo-JID, so unlike a real multi-device
// XMPP account, no more than one device id is ever needed.
const ownDeviceIDNum uint32 = 1

// OwnDeviceID is ownDeviceIDNum in e2ee.DeviceID form, for callers that
// need to reference "the" device of an owned identity (e.g. PublishedBundle).
var OwnDeviceID = e2ee.DeviceID(strconv.FormatUint(uint64(ownDeviceIDNum), 10))

// Backend is the OMEMO (XEP-0384) implementation of e2ee.Backend. It always
// supports every Version this build is capable of (see Variants) - a
// single identity/prekey set serves all of them simultaneously (XEP-0384's
// signed prekeys carry both a legacy and an OMEMO signature precisely so
// this works, see libsignal.SignedPreKeyInfo) - rather than being
// configured for one fixed version. Which version is actually used for a
// given chat's outgoing content is negotiated per chat (see
// NegotiatedVariant/negotiatedVersion), not chosen at construction time.
type Backend struct {
	libctx *libsignal.Context
	db     *badgerstore.DB

	// libsignal is not internally thread-safe (see its package doc) - one
	// mutex serializes every call into it, across every identity this
	// backend manages. Simple and correct; per-identity striping would
	// allow more concurrency but isn't needed at telegabber's message
	// volume.
	mu sync.Mutex
}

// New creates an OMEMO backend backed by db.
func New(libctx *libsignal.Context, db *badgerstore.DB) *Backend {
	return &Backend{libctx: libctx, db: db}
}

var _ e2ee.Backend = (*Backend)(nil)

func (b *Backend) Name() string { return Name }

// Close implements e2ee.Backend.
func (b *Backend) Close() error { return b.db.Close() }

// OwnDevice implements e2ee.Backend.
func (b *Backend) OwnDevice() e2ee.DeviceID { return OwnDeviceID }

func (b *Backend) Namespaces() e2ee.Namespaces {
	disco := []string{omemo0NS + ".devicelist+notify"}
	if libsignal.ProtocolV4Supported {
		disco = append(disco, omemo1NS+":devices+notify", omemo2NS+":devices+notify")
	}
	return e2ee.Namespaces{Disco: disco}
}

// Variants implements e2ee.Backend. Omemo1/Omemo2 (both libsignal.ProtocolVersionV4)
// are only listed when this build's crypto library actually supports
// protocol v4 (see libsignal.ProtocolV4Supported - false for the
// signal_legacy build tag, vanilla libsignal-protocol-c).
func (b *Backend) Variants() []e2ee.Variant {
	variants := []e2ee.Variant{e2ee.Variant(Omemo0.String())}
	if libsignal.ProtocolV4Supported {
		variants = append(variants, e2ee.Variant(Omemo1.String()), e2ee.Variant(Omemo2.String()))
	}
	return variants
}

// DefaultVariant implements e2ee.Backend: Omemo0, the most widely deployed
// version, is always what a brand new chat starts out using.
func (b *Backend) DefaultVariant() e2ee.Variant {
	return e2ee.Variant(Omemo0.String())
}

// NegotiatedVariant implements e2ee.Backend.
//
// KNOWN LIMITATION: this only ever reflects whatever variant a chat's
// session was ESTABLISHED at (via IngestRemoteBundle, or reactively via
// Decrypt processing a peer's first PreKeyMessage) - there is currently no
// way to re-negotiate an EXISTING chat to a different variant later.
// Reading libomemo-c's session_builder.c/ratchet.c directly shows that
// session_builder_process_pre_key_bundle/process_pre_key_signal_message
// always archive the OLD session state and have the replacement INHERIT
// its version (session_record_archive_current_state literally copies
// session_record_get_version(record) onto the fresh state) whenever a
// session record already exists for that address - only a genuinely fresh
// (never-before-seen) address gets its version set from what's actually
// requested. So encrypting/decrypting with a new variant against an
// address that already has a session silently keeps using the OLD
// session's version regardless of what's asked for - it does not, and
// cannot, upgrade in place. A future enhancement wanting real
// re-negotiation would need to explicitly delete the existing session
// (Store.DeleteSession/DeleteAllSessions) before re-ingesting a bundle
// under the new variant.
func (b *Backend) NegotiatedVariant(owner e2ee.PeerID) (e2ee.Variant, error) {
	login, ownerJID, ok := e2ee.SplitOwnedPeer(owner)
	if !ok {
		return "", fmt.Errorf("omemo: NegotiatedVariant: %q is not an OwnedPeer", owner)
	}

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

	version, err := b.negotiatedVersion(b.db.Store(b.libctx, login, ownerJID))
	if err != nil {
		return "", err
	}
	return e2ee.Variant(version.String()), nil
}

// negotiatedVersion is NegotiatedVariant's internal counterpart, returning
// a Version directly for Encrypt's own use - store must already be scoped
// to the right (login, owner) and b.mu must already be held.
func (b *Backend) negotiatedVersion(store *badgerstore.Store) (Version, error) {
	raw, ok, err := store.NegotiatedVersion()
	if err != nil {
		return 0, err
	}
	if !ok {
		return Omemo0, nil
	}
	return Version(raw), nil
}

// DeviceListNode implements e2ee.Backend. Node names follow each version's
// own established convention: Omemo0 uses siacs' original
// "<ns>.devicelist"; Omemo1/Omemo2 use XEP-0384's "<ns>:devices" (confirmed
// against the archived 0.4.0 spec text directly, not assumed).
func (b *Backend) DeviceListNode(variant e2ee.Variant) string {
	version, err := ParseVersion(string(variant))
	if err != nil {
		return ""
	}
	switch version {
	case Omemo0:
		return omemo0NS + ".devicelist"
	case Omemo1:
		return omemo1NS + ":devices"
	default:
		return omemo2NS + ":devices"
	}
}

// BundleNode implements e2ee.Backend. Omemo0 addresses each device's bundle
// under its own node ("<ns>.bundles:<device>"); Omemo1/Omemo2 instead use
// one shared node per XEP-0384 0.4.0's bundle-storage redesign ("Each
// bundle MUST be stored in a se[p]arate item. The item id MUST be set to
// the device id.") - confirmed against the archived 0.4.0 spec text, not
// assumed. Either way, e2ee/fetch.go only ever needs the most recent item
// on whatever node this returns, so callers don't need to know which shape
// applies.
func (b *Backend) BundleNode(variant e2ee.Variant, device e2ee.DeviceID) string {
	version, err := ParseVersion(string(variant))
	if err != nil {
		return ""
	}
	if version == Omemo0 {
		return fmt.Sprintf("%s.bundles:%s", omemo0NS, device)
	}
	ns := omemo1NS
	if version == Omemo2 {
		ns = omemo2NS
	}
	return ns + ":bundles"
}

func (b *Backend) EnsureIdentity(peer e2ee.PeerID) error {
	login, owner, ok := e2ee.SplitOwnedPeer(peer)
	if !ok {
		return fmt.Errorf("omemo: EnsureIdentity: %q is not an OwnedPeer", peer)
	}

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

	return b.ensureIdentity(b.db.Store(b.libctx, login, owner))
}

// ensureIdentity is EnsureIdentity's store-scoped counterpart, called
// defensively at the top of every other method that touches identity/
// prekey material (PublishedIdentity, PublishedBundle, Encrypt, Decrypt) -
// nothing in this package used to call the public EnsureIdentity at all,
// so every one of those would have failed with "EnsureIdentity was not
// called" the first time a chat went OMEMO-active, whichever entry point
// hit first (an inbound decrypt, an outbound encrypt, or a peer's PEP GET
// arriving before either). Bootstrapping here instead of trusting callers
// to remember makes that invariant impossible to violate. store must
// already be scoped to the right (login, owner) and b.mu must already be
// held.
func (b *Backend) ensureIdentity(store *badgerstore.Store) error {
	has, err := store.HasIdentityKeyPair()
	if err != nil {
		return fmt.Errorf("omemo: HasIdentityKeyPair: %w", err)
	}
	if has {
		return nil // idempotent - already set up
	}

	idKeyPair, err := libsignal.GenerateIdentityKeyPair(b.libctx)
	if err != nil {
		return fmt.Errorf("omemo: GenerateIdentityKeyPair: %w", err)
	}
	if err := store.SaveIdentityKeyPair(idKeyPair.Record); err != nil {
		return fmt.Errorf("omemo: SaveIdentityKeyPair: %w", err)
	}

	regID, err := libsignal.GenerateRegistrationID(b.libctx)
	if err != nil {
		return fmt.Errorf("omemo: GenerateRegistrationID: %w", err)
	}
	if err := store.SaveRegistrationID(regID); err != nil {
		return fmt.Errorf("omemo: SaveRegistrationID: %w", err)
	}

	signedPreKey, err := libsignal.GenerateSignedPreKey(b.libctx, idKeyPair, 1)
	if err != nil {
		return fmt.Errorf("omemo: GenerateSignedPreKey: %w", err)
	}
	if err := store.StoreSignedPreKey(signedPreKey.ID, signedPreKey.Record); err != nil {
		return fmt.Errorf("omemo: StoreSignedPreKey: %w", err)
	}

	start, err := store.NextPreKeyIDs(initialPreKeyCount)
	if err != nil {
		return fmt.Errorf("omemo: NextPreKeyIDs: %w", err)
	}
	preKeys, err := libsignal.GeneratePreKeys(b.libctx, start, initialPreKeyCount)
	if err != nil {
		return fmt.Errorf("omemo: GeneratePreKeys: %w", err)
	}
	for _, pk := range preKeys {
		if err := store.StorePreKey(pk.ID, pk.Record); err != nil {
			return fmt.Errorf("omemo: StorePreKey %d: %w", pk.ID, err)
		}
	}

	return nil
}

func (b *Backend) PublishedIdentity(peer e2ee.PeerID, variant e2ee.Variant) (e2ee.DeviceListDoc, error) {
	login, owner, ok := e2ee.SplitOwnedPeer(peer)
	if !ok {
		return e2ee.DeviceListDoc{}, fmt.Errorf("omemo: PublishedIdentity: %q is not an OwnedPeer", peer)
	}
	version, err := ParseVersion(string(variant))
	if err != nil {
		return e2ee.DeviceListDoc{}, fmt.Errorf("omemo: PublishedIdentity: %w", err)
	}

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

	store := b.db.Store(b.libctx, login, owner)
	if err := b.ensureIdentity(store); err != nil {
		return e2ee.DeviceListDoc{}, err
	}

	raw, err := encodeDeviceList(version, ownDeviceIDNum)
	if err != nil {
		return e2ee.DeviceListDoc{}, err
	}
	return e2ee.DeviceListDoc{Raw: raw}, nil
}

func (b *Backend) PublishedBundle(peer e2ee.PeerID, device e2ee.DeviceID, variant e2ee.Variant) (e2ee.BundleDoc, error) {
	login, owner, ok := e2ee.SplitOwnedPeer(peer)
	if !ok {
		return e2ee.BundleDoc{}, fmt.Errorf("omemo: PublishedBundle: %q is not an OwnedPeer", peer)
	}
	if device != OwnDeviceID {
		return e2ee.BundleDoc{}, fmt.Errorf("omemo: PublishedBundle: unknown device %q", device)
	}
	version, err := ParseVersion(string(variant))
	if err != nil {
		return e2ee.BundleDoc{}, fmt.Errorf("omemo: PublishedBundle: %w", err)
	}

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

	store := b.db.Store(b.libctx, login, owner)
	if err := b.ensureIdentity(store); err != nil {
		return e2ee.BundleDoc{}, err
	}

	identityPublic, _, err := store.GetIdentityKeyPair()
	if err != nil {
		return e2ee.BundleDoc{}, fmt.Errorf("omemo: GetIdentityKeyPair: %w", err)
	}

	// ensureIdentity always creates signed prekey id 1 and never rotates it
	// yet (a future enhancement), so id 1 is the only one that will exist.
	spkRecord, found, err := store.LoadSignedPreKey(1)
	if err != nil {
		return e2ee.BundleDoc{}, fmt.Errorf("omemo: LoadSignedPreKey: %w", err)
	}
	if !found {
		return e2ee.BundleDoc{}, errors.New("omemo: PublishedBundle: signed prekey 1 missing despite ensureIdentity succeeding")
	}
	spkInfo, err := libsignal.DecodeSignedPreKey(b.libctx, spkRecord)
	if err != nil {
		return e2ee.BundleDoc{}, fmt.Errorf("omemo: DecodeSignedPreKey: %w", err)
	}
	// See the rule documented on libsignal.SignedPreKeyInfo: the signature
	// must match the serialization form the verifier will re-derive, which
	// depends on the requested variant (Omemo0 -> legacy form, Omemo1/Omemo2
	// -> OMEMO form, since both share libsignal.ProtocolVersionV4) - the
	// same signed prekey serves every variant, just re-signed per request.
	signature := spkInfo.Signature
	if version != Omemo0 {
		signature = spkInfo.SignatureOMEMO
	}

	preKeyIDs, err := store.ListPreKeyIDs()
	if err != nil {
		return e2ee.BundleDoc{}, fmt.Errorf("omemo: ListPreKeyIDs: %w", err)
	}
	preKeys := make([]libsignal.PreKeyInfo, 0, len(preKeyIDs))
	for _, id := range preKeyIDs {
		record, found, err := store.LoadPreKey(id)
		if err != nil {
			return e2ee.BundleDoc{}, fmt.Errorf("omemo: LoadPreKey %d: %w", id, err)
		}
		if !found {
			continue // consumed between ListPreKeyIDs and here - fine, just skip it
		}
		info, err := libsignal.DecodePreKey(b.libctx, record)
		if err != nil {
			return e2ee.BundleDoc{}, fmt.Errorf("omemo: DecodePreKey %d: %w", id, err)
		}
		preKeys = append(preKeys, *info)
	}

	raw, err := encodeBundle(version, identityPublic, spkInfo.ID, spkInfo.PublicKey, signature, preKeys)
	if err != nil {
		return e2ee.BundleDoc{}, err
	}
	return e2ee.BundleDoc{Raw: raw}, nil
}

func (b *Backend) IngestRemoteDeviceList(owner, peer e2ee.PeerID, doc e2ee.DeviceListDoc) ([]e2ee.DeviceID, error) {
	login, ownerJID, ok := e2ee.SplitOwnedPeer(owner)
	if !ok {
		return nil, fmt.Errorf("omemo: IngestRemoteDeviceList: %q is not an OwnedPeer", owner)
	}
	ids, err := parseDeviceList(doc.Raw)
	if err != nil {
		return nil, fmt.Errorf("omemo: parseDeviceList: %w", err)
	}

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

	store := b.db.Store(b.libctx, login, ownerJID)
	if err := store.SaveRemoteDeviceListCache(string(peer), doc.Raw); err != nil {
		return nil, err
	}

	deviceIDs := make([]e2ee.DeviceID, len(ids))
	for i, id := range ids {
		deviceIDs[i] = e2ee.DeviceID(strconv.FormatUint(uint64(id), 10))
	}
	return deviceIDs, nil
}

func (b *Backend) IngestRemoteBundle(owner, peer e2ee.PeerID, device e2ee.DeviceID, doc e2ee.BundleDoc) error {
	login, ownerJID, ok := e2ee.SplitOwnedPeer(owner)
	if !ok {
		return fmt.Errorf("omemo: IngestRemoteBundle: %q is not an OwnedPeer", owner)
	}
	deviceID, err := parseDeviceID(device)
	if err != nil {
		return err
	}
	parsed, err := parseBundle(doc.Raw)
	if err != nil {
		return fmt.Errorf("omemo: parseBundle: %w", err)
	}

	var preKeyID uint32
	var preKeyPublic []byte
	if len(parsed.PreKeys) > 0 {
		idx, err := cryptoRandIndex(len(parsed.PreKeys))
		if err != nil {
			return err
		}
		preKeyID = parsed.PreKeys[idx].ID
		preKeyPublic = parsed.PreKeys[idx].Public
	}

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

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

	// The session's protocol version must match what the fetched bundle
	// itself actually is (parsed.Version), derived from doc.Raw's own shape
	// rather than assumed from whatever variant fetch.go asked for -
	// picking the wrong one here fails bundle signature verification with
	// SG_ERR_INVALID_KEY (the same class of bug already hit and fixed once
	// for PublishedBundle/Decrypt - see their comments).
	builder, err := libsignal.NewSessionBuilder(b.libctx, storeCtx, libsignal.Address{Name: string(peer), DeviceID: deviceID}, parsed.Version.libsignalVersion())
	if err != nil {
		return fmt.Errorf("omemo: NewSessionBuilder: %w", err)
	}
	defer builder.Close()

	err = builder.ProcessPreKeyBundle(b.libctx, libsignal.RemoteBundle{
		// OMEMO's wire format (neither legacy nor OMEMOKeyExchange) ever
		// transmits a registration id - device id already does what it
		// would otherwise disambiguate. 0 is standard, harmless practice.
		RegistrationID:        0,
		DeviceID:              deviceID,
		PreKeyID:              preKeyID,
		PreKeyPublic:          preKeyPublic,
		SignedPreKeyID:        parsed.SignedPreKeyID,
		SignedPreKeyPublic:    parsed.SignedPreKeyPublic,
		SignedPreKeySignature: parsed.SignedPreKeySignature,
		IdentityKeyPublic:     parsed.IdentityKeyPublic,
	})
	if err != nil {
		return fmt.Errorf("omemo: ProcessPreKeyBundle: %w", err)
	}

	// Keep the chat's negotiated version in sync with whatever session we
	// just established/replaced - libomemo-c's session_cipher_encrypt
	// always serializes using the session's own stored version
	// (session_state_get_session_version), NOT whatever version Encrypt
	// asks for, so a stale negotiated-version record here would make
	// Encrypt produce a WireEnvelope whose declared Version doesn't match
	// the ACTUAL bytes libsignal ends up emitting - confirmed by reading
	// session_cipher.c's session_cipher_encrypt/decrypt directly, not
	// assumed (this is exactly the SG_ERR_INVALID_PROTO_BUF failure the
	// old, purely Decrypt-driven negotiation tracking could hit as soon as
	// a chat's session got established/replaced at a version other than
	// its previously-recorded negotiated one).
	if err := store.SetNegotiatedVersion(byte(parsed.Version)); err != nil {
		return fmt.Errorf("omemo: SetNegotiatedVersion: %w", err)
	}
	return nil
}

func (b *Backend) Devices(owner, peer e2ee.PeerID) ([]e2ee.DeviceInfo, error) {
	login, ownerJID, ok := e2ee.SplitOwnedPeer(owner)
	if !ok {
		return nil, fmt.Errorf("omemo: Devices: %q is not an OwnedPeer", owner)
	}

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

	store := b.db.Store(b.libctx, login, ownerJID)
	deviceIDs, err := store.GetSubDeviceSessions(string(peer))
	if err != nil {
		return nil, err
	}

	infos := make([]e2ee.DeviceInfo, 0, len(deviceIDs))
	for _, id := range deviceIDs {
		// A stored session implies IsTrustedIdentity already accepted this
		// device's key (TOFU) - finer-grained trust states (manually
		// verified/revoked) are a documented future enhancement, not
		// needed for this phase.
		infos = append(infos, e2ee.DeviceInfo{ID: e2ee.DeviceID(strconv.FormatUint(uint64(id), 10)), Trusted: true})
	}
	return infos, nil
}

func (b *Backend) Enabled(owner e2ee.PeerID) (bool, error) {
	login, ownerJID, ok := e2ee.SplitOwnedPeer(owner)
	if !ok {
		return false, fmt.Errorf("omemo: Enabled: %q is not an OwnedPeer", owner)
	}

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

	return b.db.Store(b.libctx, login, ownerJID).Enabled()
}

func (b *Backend) SetEnabled(owner e2ee.PeerID, enabled bool) error {
	login, ownerJID, ok := e2ee.SplitOwnedPeer(owner)
	if !ok {
		return fmt.Errorf("omemo: SetEnabled: %q is not an OwnedPeer", owner)
	}

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

	return b.db.Store(b.libctx, login, ownerJID).SetEnabled(enabled)
}

func parseDeviceID(device e2ee.DeviceID) (uint32, error) {
	id, err := strconv.ParseUint(string(device), 10, 32)
	if err != nil {
		return 0, fmt.Errorf("omemo: invalid device id %q: %w", device, err)
	}
	return uint32(id), nil
}

func cryptoRandIndex(n int) (int, error) {
	i, err := rand.Int(rand.Reader, big.NewInt(int64(n)))
	if err != nil {
		return 0, err
	}
	return int(i.Int64()), nil
}