package membership

import (
	"bytes"
	"path/filepath"
	"testing"
)

func openTestStore(t *testing.T) *Store {
	t.Helper()
	path := filepath.Join(t.TempDir(), "test.db")
	s, err := Open(path)
	if err != nil {
		t.Fatalf("Open: %v", err)
	}
	t.Cleanup(func() { s.Close() })
	return s
}

func TestMigrationsCreateSchema(t *testing.T) {
	s := openTestStore(t)
	// Verify the three tables exist by querying sqlite_master.
	for _, tbl := range []string{"rooms", "members", "room_keys"} {
		var name string
		err := s.db.QueryRow(
			`SELECT name FROM sqlite_master WHERE type='table' AND name=?`, tbl,
		).Scan(&name)
		if err != nil {
			t.Fatalf("table %q not created: %v", tbl, err)
		}
	}
	// Re-applying migrations must be idempotent (no error on a populated db).
	if err := s.applyMigrations(); err != nil {
		t.Fatalf("re-apply migrations: %v", err)
	}
}

func TestListRoomsForEndpoint(t *testing.T) {
	s := openTestStore(t)

	// Owner of two rooms; a member in only the first.
	owner, member := "owner-ep", "member-ep"
	mk := func(id, subj string) RoomInfo {
		return RoomInfo{RoomID: id, Subject: subj, Encrypt: true, Persist: true, SignMsgs: true, OwnerEndpoint: owner}
	}
	if err := s.CreateRoom(mk("room-a", "room.a"), []byte("os"), []byte("ok"), []byte("k")); err != nil {
		t.Fatalf("CreateRoom a: %v", err)
	}
	if err := s.CreateRoom(mk("room-b", "room.b"), []byte("os"), []byte("ok"), []byte("k")); err != nil {
		t.Fatalf("CreateRoom b: %v", err)
	}
	if err := s.AddMember("room-a", Member{Endpoint: member, Role: "member", SignPub: []byte("s"), KexPub: []byte("k")}, 1, []byte("mk")); err != nil {
		t.Fatalf("AddMember: %v", err)
	}

	// Owner is in both rooms, as owner, ordered by room id.
	ownerRooms, err := s.ListRoomsForEndpoint(owner)
	if err != nil {
		t.Fatalf("ListRoomsForEndpoint owner: %v", err)
	}
	if len(ownerRooms) != 2 {
		t.Fatalf("owner: expected 2 rooms, got %d", len(ownerRooms))
	}
	if ownerRooms[0].RoomID != "room-a" || ownerRooms[1].RoomID != "room-b" {
		t.Fatalf("owner rooms not ordered: %+v", ownerRooms)
	}
	if ownerRooms[0].Role != "owner" || !ownerRooms[0].Encrypt || ownerRooms[0].Subject != "room.a" {
		t.Fatalf("owner room metadata wrong: %+v", ownerRooms[0])
	}

	// Member is in exactly one room, as member.
	memberRooms, err := s.ListRoomsForEndpoint(member)
	if err != nil {
		t.Fatalf("ListRoomsForEndpoint member: %v", err)
	}
	if len(memberRooms) != 1 || memberRooms[0].RoomID != "room-a" || memberRooms[0].Role != "member" {
		t.Fatalf("member rooms wrong: %+v", memberRooms)
	}

	// An unknown endpoint yields an empty slice, not an error.
	none, err := s.ListRoomsForEndpoint("nobody")
	if err != nil {
		t.Fatalf("ListRoomsForEndpoint nobody: %v", err)
	}
	if len(none) != 0 {
		t.Fatalf("expected no rooms for unknown endpoint, got %+v", none)
	}
}

func TestRoomMemberKeyRoundTrip(t *testing.T) {
	s := openTestStore(t)

	owner := "owner-ep"
	roomID := "room-1"
	info := RoomInfo{
		RoomID:        roomID,
		Subject:       "room.test",
		Encrypt:       true,
		Persist:       true,
		SignMsgs:      true,
		OwnerEndpoint: owner,
	}
	ownerSealed := []byte("owner-sealed-key-epoch1")
	if err := s.CreateRoom(info, []byte("owner-sign"), []byte("owner-kex"), ownerSealed); err != nil {
		t.Fatalf("CreateRoom: %v", err)
	}

	// GetRoom returns epoch 1 and the policy.
	got, err := s.GetRoom(roomID)
	if err != nil {
		t.Fatalf("GetRoom: %v", err)
	}
	if got.Epoch != 1 || !got.Encrypt || !got.Persist || !got.SignMsgs || got.OwnerEndpoint != owner {
		t.Fatalf("GetRoom mismatch: %+v", got)
	}

	// Owner sealed key at epoch 1 (latest).
	ep, sealed, err := s.GetSealedKey(roomID, owner, 0)
	if err != nil {
		t.Fatalf("GetSealedKey owner: %v", err)
	}
	if ep != 1 || !bytes.Equal(sealed, ownerSealed) {
		t.Fatalf("owner sealed key mismatch: epoch=%d sealed=%q", ep, sealed)
	}

	// Add member at epoch 1.
	member := Member{Endpoint: "member-ep", Role: "member", SignPub: []byte("m-sign"), KexPub: []byte("m-kex")}
	memberSealed := []byte("member-sealed-epoch1")
	if err := s.AddMember(roomID, member, 1, memberSealed); err != nil {
		t.Fatalf("AddMember: %v", err)
	}

	gotMember, err := s.GetMember(roomID, "member-ep")
	if err != nil {
		t.Fatalf("GetMember: %v", err)
	}
	if gotMember.Role != "member" || !bytes.Equal(gotMember.SignPub, []byte("m-sign")) {
		t.Fatalf("GetMember mismatch: %+v", gotMember)
	}

	members, err := s.ListMembers(roomID)
	if err != nil {
		t.Fatalf("ListMembers: %v", err)
	}
	if len(members) != 2 {
		t.Fatalf("expected 2 members, got %d", len(members))
	}

	// Bump to epoch 2 and store new keys only for the owner (simulating a kick of member-ep).
	if err := s.BumpEpoch(roomID, 2); err != nil {
		t.Fatalf("BumpEpoch: %v", err)
	}
	newKeys := map[string][]byte{owner: []byte("owner-sealed-epoch2")}
	if err := s.PutSealedKeys(roomID, 2, newKeys); err != nil {
		t.Fatalf("PutSealedKeys: %v", err)
	}
	if err := s.RemoveMember(roomID, "member-ep"); err != nil {
		t.Fatalf("RemoveMember: %v", err)
	}

	got, err = s.GetRoom(roomID)
	if err != nil {
		t.Fatalf("GetRoom after bump: %v", err)
	}
	if got.Epoch != 2 {
		t.Fatalf("expected epoch 2, got %d", got.Epoch)
	}

	// Owner now has a fresh sealed key at epoch 2 (latest).
	ep, sealed, err = s.GetSealedKey(roomID, owner, 0)
	if err != nil {
		t.Fatalf("GetSealedKey owner epoch2: %v", err)
	}
	if ep != 2 || !bytes.Equal(sealed, []byte("owner-sealed-epoch2")) {
		t.Fatalf("owner epoch2 key mismatch: epoch=%d sealed=%q", ep, sealed)
	}

	// The removed member is gone.
	if _, err := s.GetMember(roomID, "member-ep"); err == nil {
		t.Fatalf("expected error getting removed member")
	}

	// The kicked member has no key at epoch 2 (was excluded from the rekey).
	if _, _, err := s.GetSealedKey(roomID, "member-ep", 2); err == nil {
		t.Fatalf("kicked member should have no key at epoch 2")
	}
	members, err = s.ListMembers(roomID)
	if err != nil {
		t.Fatalf("ListMembers after remove: %v", err)
	}
	if len(members) != 1 || members[0].Endpoint != owner {
		t.Fatalf("expected only owner remaining, got %+v", members)
	}
}