package syncapi

import (
	"context"
	"slices"

	"github.com/dariuszjaniak/warp/server/internal/matrix/event"
	"github.com/dariuszjaniak/warp/server/internal/roomserver"
)

// Device list tracking, the reading side: whose end-to-end keys a client
// should fetch again, and whose it can forget.
//
// A client encrypting for a room needs the keys of every device of every
// member, and caches them. The cache goes stale in two ways, and the two
// lists of a sync's device_lists are one each:
//
//   - changed: a user who shares a room with the client has a new device,
//     lost one, or changed their cross-signing keys — or has only just come
//     to share a room with the client, whose cache of their keys, if it has
//     one at all, is from before.
//   - left: a user who no longer shares any room with the client, whose
//     keys it can stop tracking.
//
// The first kind comes from the keyserver's device list stream, which
// moves for local users when they upload or delete keys and for remote
// users when their server says so. The second kind comes from membership:
// whose membership moved in a room this user is or was in, and whether
// the two still share a room once it did. Both are read only on an
// incremental sync; an initial one is the client's cue to fetch every key
// it needs anyway.

// deviceLists works out the device_lists of a sync between since and to.
func (s *API) deviceLists(ctx context.Context, userID string, since, to Token, rooms []roomserver.MembershipAt) (DeviceLists, error) {
	out := DeviceLists{Changed: []string{}, Left: []string{}}
	if s.keys == nil || since.IsInitial() {
		return out, nil
	}
	keysMoved := to.DeviceList > since.DeviceList
	roomsMoved := to.PDU > since.PDU
	if !keysMoved && !roomsMoved {
		return out, nil
	}

	sharing, err := s.rooms.UsersSharingRooms(ctx, userID)
	if err != nil {
		return DeviceLists{}, err
	}
	shares := func(u string) bool { return slices.Contains(sharing, u) }

	changed := map[string]bool{}
	left := map[string]bool{}

	if keysMoved {
		// The user's own devices are reported too: another of their
		// devices has to hear about the new one to start encrypting for
		// it.
		watched := append(slices.Clone(sharing), userID)
		moves, err := s.keys.DeviceListChanges(ctx, watched, since.DeviceList, to.DeviceList)
		if err != nil {
			return DeviceLists{}, err
		}
		for _, m := range moves {
			changed[m.UserID] = true
		}
	}

	if roomsMoved {
		if err := s.membershipDeviceLists(ctx, userID, since, to, rooms, shares, changed, left); err != nil {
			return DeviceLists{}, err
		}
	}

	for u := range changed {
		out.Changed = append(out.Changed, u)
		delete(left, u)
	}
	for u := range left {
		out.Left = append(out.Left, u)
	}
	slices.Sort(out.Changed)
	slices.Sort(out.Left)
	return out, nil
}

// membershipDeviceLists adds the users whose sharing of a room with userID
// began or ended inside the window.
func (s *API) membershipDeviceLists(
	ctx context.Context,
	userID string,
	since, to Token,
	rooms []roomserver.MembershipAt,
	shares func(string) bool,
	changed, left map[string]bool,
) error {
	consider := func(other string) {
		switch {
		case other == userID:
		case shares(other):
			changed[other] = true
		default:
			left[other] = true
		}
	}

	// A room whose membership of this user moved in the window: joined,
	// everybody in it is new to share with; left, everybody in it may no
	// longer be.
	for _, room := range rooms {
		if room.Position <= since.PDU || room.Position > to.PDU {
			continue
		}
		if room.Membership != event.MembershipJoin && !isDeparted(room.Membership) {
			continue
		}
		members, err := s.rooms.JoinedUserIDs(ctx, room.RoomID)
		if err != nil {
			return err
		}
		for _, member := range members {
			consider(member)
		}
	}

	// Everybody else whose membership moved in a room this user is or was
	// in. An invite or a knock is not a change of who shares the room, and
	// is skipped: telling a client to forget the keys of a user it was
	// encrypting for as an invitee would have it stop.
	changes, err := s.rooms.MemberChanges(ctx, roomIDsOf(rooms), since.PDU, to.PDU)
	if err != nil {
		return err
	}
	for _, c := range changes {
		if c.Membership == event.MembershipJoin || isDeparted(c.Membership) {
			consider(c.UserID)
		}
	}
	return nil
}

// KeyChanges answers GET /keys/changes: whose device lists changed between
// two tokens, for a client that missed some syncs, or is catching up with
// the history of a room it was not syncing.
func (s *API) KeyChanges(ctx context.Context, userID, from, to string) (DeviceLists, error) {
	since, err := ParseToken(from)
	if err != nil {
		return DeviceLists{}, errInvalidTokenf("from: %v", err)
	}
	var until Token
	if to == "" {
		positions, err := s.store.StreamPositions(ctx)
		if err != nil {
			return DeviceLists{}, err
		}
		until = tokenAt(positions)
	} else if until, err = ParseToken(to); err != nil {
		return DeviceLists{}, errInvalidTokenf("to: %v", err)
	}

	if since.IsInitial() {
		// From the beginning of time is every device list there is, which
		// is the client fetching everybody's keys: what it would do after
		// an initial sync anyway, and what since=0 means everywhere else.
		// Reported as every user it shares a room with.
		sharing, err := s.rooms.UsersSharingRooms(ctx, userID)
		if err != nil {
			return DeviceLists{}, err
		}
		if sharing == nil {
			sharing = []string{}
		}
		return DeviceLists{Changed: sharing, Left: []string{}}, nil
	}

	rooms, err := s.rooms.UserMemberships(ctx, userID, nil)
	if err != nil {
		return DeviceLists{}, err
	}
	return s.deviceLists(ctx, userID, since, until, rooms)
}

// deviceKeyCounts reads the key counts of the syncing device.
func (s *API) deviceKeyCounts(ctx context.Context, req SyncRequest, resp *Response) error {
	if s.keys == nil || req.Device == "" {
		return nil
	}
	counts, err := s.keys.OneTimeKeyCounts(ctx, req.User, req.Device)
	if err != nil {
		return err
	}
	unused, err := s.keys.UnusedFallbackKeyTypes(ctx, req.User, req.Device)
	if err != nil {
		return err
	}
	resp.DeviceOneTimeKeysCount = counts
	resp.DeviceUnusedFallbackKeyTypes = &unused
	return nil
}