package omemo_test

import (
	"bytes"
	"testing"

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

func newTestBackend(t *testing.T, version omemo.Version) *omemo.Backend {
	t.Helper()

	ctx, err := libsignal.NewContext()
	if err != nil {
		t.Fatalf("NewContext: %v", err)
	}
	t.Cleanup(ctx.Close)

	db, err := badgerstore.Open(t.TempDir(), nil)
	if err != nil {
		t.Fatalf("badgerstore.Open: %v", err)
	}
	t.Cleanup(func() { db.Close() })

	return omemo.New(ctx, db, version)
}

// TestBackendRoundTrip exercises the full e2ee.Backend interface, not just
// the underlying libsignal primitives (already covered by libsignal's own
// round-trip test): EnsureIdentity, publishing a device-list/bundle,
// ingesting them on the other side, and Encrypt/Decrypt in both
// directions. This simulates two independent parties - "gateway" (the
// bridged chat pseudo-JID's owned identity) and "client" (the real XMPP
// user's own device) - each with their own Backend/store, exchanging
// PublishedIdentity/PublishedBundle results directly in place of the real
// PEP fetch (e2ee/fetch.go's actual IQ round trip isn't built yet).
func TestBackendRoundTrip(t *testing.T) {
	const gatewayLogin = "telegram-login-1"
	const gatewayBareJID = "12345@transport.example"
	const realUserBareJID = "alice@real.example"

	for _, tc := range []struct {
		name    string
		version omemo.Version
	}{
		{"omemo0", omemo.Omemo0},
		{"omemo1", omemo.Omemo1},
		{"omemo2", omemo.Omemo2},
	} {
		t.Run(tc.name, func(t *testing.T) {
			if tc.version != omemo.Omemo0 && !libsignal.ProtocolV4Supported {
				t.Skip("protocol v4 (omemo1/omemo2) not supported by this build (signal_legacy tag)")
			}

			gateway := newTestBackend(t, tc.version)
			client := newTestBackend(t, tc.version)

			gatewayOwned := e2ee.OwnedPeer(gatewayLogin, gatewayBareJID)
			clientOwned := e2ee.OwnedPeer("n/a", realUserBareJID)
			realUser := e2ee.PeerID(realUserBareJID)     // how the gateway refers to the real user
			gatewayAsPeer := e2ee.PeerID(gatewayBareJID) // how the client refers to the gateway

			if err := gateway.EnsureIdentity(gatewayOwned); err != nil {
				t.Fatalf("gateway.EnsureIdentity: %v", err)
			}
			if err := client.EnsureIdentity(clientOwned); err != nil {
				t.Fatalf("client.EnsureIdentity: %v", err)
			}

			// Gateway "fetches" the client's device list + bundle.
			clientDeviceList, err := client.PublishedIdentity(clientOwned)
			if err != nil {
				t.Fatalf("client.PublishedIdentity: %v", err)
			}
			if err := gateway.IngestRemoteDeviceList(gatewayOwned, realUser, clientDeviceList); err != nil {
				t.Fatalf("gateway.IngestRemoteDeviceList: %v", err)
			}
			clientBundle, err := client.PublishedBundle(clientOwned, omemo.OwnDeviceID)
			if err != nil {
				t.Fatalf("client.PublishedBundle: %v", err)
			}
			if err := gateway.IngestRemoteBundle(gatewayOwned, realUser, omemo.OwnDeviceID, clientBundle); err != nil {
				t.Fatalf("gateway.IngestRemoteBundle: %v", err)
			}

			// Gateway encrypts a message originating from Telegram.
			plaintext1 := []byte("hello from telegram")
			env1, err := gateway.Encrypt(gatewayOwned, []e2ee.PeerID{realUser}, plaintext1)
			if err != nil {
				t.Fatalf("gateway.Encrypt: %v", err)
			}

			// Client decrypts it - this establishes its side of the
			// session as a side effect, with no prior bundle fetch needed
			// on its end (matching the prekey-message responder flow
			// already validated at the libsignal level).
			got1, err := client.Decrypt(gatewayAsPeer, clientOwned, env1)
			if err != nil {
				t.Fatalf("client.Decrypt: %v", err)
			}
			if !bytes.Equal(got1, plaintext1) {
				t.Fatalf("round trip mismatch: got %q, want %q", got1, plaintext1)
			}

			// Client replies. It needs no explicit IngestRemoteBundle call
			// for the gateway's device - the session established during
			// the decrypt above already covers it.
			plaintext2 := []byte("hi telegram")
			env2, err := client.Encrypt(clientOwned, []e2ee.PeerID{gatewayAsPeer}, plaintext2)
			if err != nil {
				t.Fatalf("client.Encrypt: %v", err)
			}

			got2, err := gateway.Decrypt(realUser, gatewayOwned, env2)
			if err != nil {
				t.Fatalf("gateway.Decrypt: %v", err)
			}
			if !bytes.Equal(got2, plaintext2) {
				t.Fatalf("round trip mismatch: got %q, want %q", got2, plaintext2)
			}
		})
	}
}