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) *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)
}

// 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, Encrypt/Decrypt in both directions, and
// per-chat version negotiation (defaults to Omemo0, then tracks whatever
// variant the chat's very first session actually gets established at - see
// e2ee.Backend.NegotiatedVariant's doc comment). Each variant is tested with
// its own fresh pair of identities/sessions rather than trying to "upgrade"
// one shared chat mid-test - see testBackendRoundTripAtVariant's own doc
// comment for why that's the only valid way to exercise a non-default
// variant with this crypto library.
func TestBackendRoundTrip(t *testing.T) {
	t.Run("omemo0 (default)", func(t *testing.T) {
		testBackendRoundTripAtVariant(t, "telegram-login-1", "12345@transport.example", "alice@real.example", e2ee.Variant(omemo.Omemo0.String()))
	})

	t.Run("omemo2 (fresh chat established directly at a non-default variant)", func(t *testing.T) {
		if !libsignal.ProtocolV4Supported {
			t.Skip("protocol v4 (omemo1/omemo2) not supported by this build (signal_legacy tag)")
		}
		testBackendRoundTripAtVariant(t, "telegram-login-2", "67890@transport.example", "bob@real.example", e2ee.Variant(omemo.Omemo2.String()))
	})
}

// testBackendRoundTripAtVariant simulates two independent parties -
// "gateway" (the bridged chat pseudo-JID's owned identity) and "client"
// (the real XMPP user's own device) - exchanging PublishedIdentity/
// PublishedBundle results directly in place of the real PEP fetch
// (e2ee/fetch.go's actual IQ round trip isn't exercised here, only
// Backend's own contract), with their very first session established at
// variant.
//
// This can only test a variant from a chat's FIRST-EVER message, not by
// switching an EXISTING session to a different variant mid-test: reading
// libomemo-c's session_builder.c and 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 new one INHERIT its
// version (session_record_archive_current_state calls
// session_state_set_session_version(new_state, session_record_get_version(record))
// using the OLD record's version) 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
// (signal_protocol_session_load_session's "record didn't exist yet"
// branch calls session_record_set_version with the requested version).
// Concretely: encrypting/decrypting with a NEW protocol version against an
// address that already has a session silently keeps using the OLD
// session's version regardless of what's requested - this is what
// produced a "SessionCipher.Decrypt: ... invalid protobuf" failure the
// first time this test tried exactly that. See the matching note on
// e2ee.Backend.NegotiatedVariant: telegabber does not yet support
// re-negotiating an EXISTING chat to a different variant (that would
// require explicitly deleting the old session first) - only the
// auto-upgrade-from-nothing case (this test) is implemented.
func testBackendRoundTripAtVariant(t *testing.T, gatewayLogin, gatewayBareJID, realUserBareJID string, variant e2ee.Variant) {
	t.Helper()

	gateway := newTestBackend(t)
	client := newTestBackend(t)

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

	if v, err := gateway.NegotiatedVariant(gatewayOwned); err != nil || v != gateway.DefaultVariant() {
		t.Fatalf("expected a fresh chat to default to %v, got %v (err %v)", gateway.DefaultVariant(), v, err)
	}

	// Gateway "fetches" the client's device list + bundle, under variant.
	clientDeviceList, err := client.PublishedIdentity(clientOwned, variant)
	if err != nil {
		t.Fatalf("client.PublishedIdentity: %v", err)
	}
	gotDeviceIDs, err := gateway.IngestRemoteDeviceList(gatewayOwned, realUser, clientDeviceList)
	if err != nil {
		t.Fatalf("gateway.IngestRemoteDeviceList: %v", err)
	}
	if len(gotDeviceIDs) != 1 || gotDeviceIDs[0] != omemo.OwnDeviceID {
		t.Fatalf("expected device list to contain only %v, got %v", omemo.OwnDeviceID, gotDeviceIDs)
	}
	clientBundle, err := client.PublishedBundle(clientOwned, omemo.OwnDeviceID, variant)
	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)
	}
	if v, err := gateway.NegotiatedVariant(gatewayOwned); err != nil || v != variant {
		t.Fatalf("expected gateway to have negotiated %v after ingesting its bundle, got %v (err %v)", variant, 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), and records variant as negotiated too.
	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)
	}
	if v, err := client.NegotiatedVariant(clientOwned); err != nil || v != variant {
		t.Fatalf("expected client to have negotiated %v after decrypting its first message, got %v (err %v)", variant, v, err)
	}

	// 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)
	}
}