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.
type Backend struct {
	libctx  *libsignal.Context
	db      *badgerstore.DB
	version Version // which version this backend produces for new outgoing content; incoming messages of any of the three versions are always accepted, decoded per-message

	// 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, producing new outgoing
// content in the given version (Omemo0, Omemo1, or Omemo2 - see version.go).
func New(libctx *libsignal.Context, db *badgerstore.DB, version Version) *Backend {
	return &Backend{libctx: libctx, db: db, version: version}
}

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

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

func (b *Backend) Namespaces() e2ee.Namespaces {
	switch b.version {
	case Omemo0:
		return e2ee.Namespaces{
			Disco: []string{omemo0NS + ".devicelist+notify"},
			EME:   omemo0NS,
		}
	case Omemo1:
		return e2ee.Namespaces{
			Disco: []string{omemo1NS + ":devices+notify"},
			EME:   omemo1NS,
		}
	default:
		return e2ee.Namespaces{
			Disco: []string{omemo2NS + ":devices+notify"},
			EME:   omemo2NS,
		}
	}
}

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()

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

	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) (e2ee.DeviceListDoc, error) {
	login, owner, ok := e2ee.SplitOwnedPeer(peer)
	if !ok {
		return e2ee.DeviceListDoc{}, fmt.Errorf("omemo: PublishedIdentity: %q is not an OwnedPeer", peer)
	}

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

	store := b.db.Store(b.libctx, login, owner)
	if has, err := store.HasIdentityKeyPair(); err != nil {
		return e2ee.DeviceListDoc{}, err
	} else if !has {
		return e2ee.DeviceListDoc{}, errors.New("omemo: PublishedIdentity: EnsureIdentity was not called")
	}

	raw, err := encodeDeviceList(b.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) (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)
	}

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

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

	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: EnsureIdentity was not called")
	}
	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 recipient's protocol version (Omemo0 -> legacy form,
	// Omemo1/Omemo2 -> OMEMO form, since both share libsignal.ProtocolVersionV4)
	// - but since we only publish one bundle version at a time (this
	// backend's own configured version), that's the one to use here too.
	signature := spkInfo.Signature
	if b.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(b.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) error {
	login, ownerJID, ok := e2ee.SplitOwnedPeer(owner)
	if !ok {
		return fmt.Errorf("omemo: IngestRemoteDeviceList: %q is not an OwnedPeer", owner)
	}
	if _, err := parseDeviceList(doc.Raw); err != nil {
		return fmt.Errorf("omemo: parseDeviceList: %w", err)
	}

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

	store := b.db.Store(b.libctx, login, ownerJID)
	return store.SaveRemoteDeviceListCache(string(peer), doc.Raw)
}

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()

	builder, err := libsignal.NewSessionBuilder(b.libctx, storeCtx, libsignal.Address{Name: string(peer), DeviceID: deviceID}, b.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)
	}
	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 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
}