package stateres

import (
	"container/heap"
	"context"
	"math"
	"slices"
	"strings"

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

// reverseTopologicalPowerSort orders events earliest to latest, so that
// nothing comes before an event that authorises it, and where there is a
// choice the sender with the most power goes first.
//
// It is Kahn's algorithm with a priority queue: at every step, of the
// events whose auth events have all been placed, the one that sorts
// smallest is taken. Tie-breaking on timestamp and then on event ID makes
// the result identical on every server, which is the whole requirement —
// two servers that ordered these differently would resolve the state
// differently and the room would fork.
func (r *resolver) reverseTopologicalPowerSort(ctx context.Context, events map[string]*event.Event) ([]*event.Event, error) {
	// Edges run from an auth event to the event it authorises, counted
	// only within the set being sorted: a reference out of the set has
	// already happened as far as this ordering goes.
	authorises := make(map[string][]string, len(events))
	remaining := make(map[string]int, len(events))
	for id, e := range events {
		for _, authID := range e.AuthEventIDs() {
			if _, inSet := events[authID]; !inSet {
				continue
			}
			authorises[authID] = append(authorises[authID], id)
			remaining[id]++
		}
	}

	keys := make(map[string]sortKey, len(events))
	for id, e := range events {
		level, err := r.senderPowerLevel(ctx, e)
		if err != nil {
			return nil, err
		}
		keys[id] = sortKey{rank: level, ts: e.OriginServerTS(), id: id}
	}

	ready := &sortHeap{keys: keys}
	for id := range events {
		if remaining[id] == 0 {
			ready.ids = append(ready.ids, id)
		}
	}
	heap.Init(ready)

	out := make([]*event.Event, 0, len(events))
	for ready.Len() > 0 {
		id := heap.Pop(ready).(string)
		out = append(out, events[id])
		for _, next := range authorises[id] {
			remaining[next]--
			if remaining[next] == 0 {
				heap.Push(ready, next)
			}
		}
	}

	// A cycle in the auth graph would leave events unplaced. It cannot
	// happen in a well-formed room, where an event's auth events always
	// precede it, but a malicious server can send anything and quietly
	// dropping state events would be worse than an arbitrary but stable
	// order.
	if len(out) < len(events) {
		leftover := make([]*event.Event, 0, len(events)-len(out))
		for id, e := range events {
			if remaining[id] > 0 {
				leftover = append(leftover, e)
			}
		}
		sortByKey(leftover, keys)
		out = append(out, leftover...)
	}
	return out, nil
}

// mainlineSort orders events by where they hang off the room's succession
// of power levels events.
//
// The mainline is that succession: the given power levels event, the one
// in its auth events, and so on back to the first one in the room. An
// event's position is how far back along it the event depends, and an
// event depending on an older link sorts first. The effect is that a
// change made under an older set of power levels is applied before one
// made under a newer set, so the newer set has the last word.
func (r *resolver) mainlineSort(ctx context.Context, events map[string]*event.Event, powerLevels *event.Event) ([]*event.Event, error) {
	// index maps each power levels event in the mainline to its place,
	// 0 being the most recent.
	index := map[string]int{}
	for current := powerLevels; current != nil; {
		if _, loop := index[current.EventID()]; loop {
			break // a malformed room; stop rather than walk for ever
		}
		index[current.EventID()] = len(index)
		next, err := r.powerLevelsInAuthEvents(ctx, current)
		if err != nil {
			return nil, err
		}
		current = next
	}

	keys := make(map[string]sortKey, len(events))
	positions := map[string]int{}
	for id, e := range events {
		position, err := r.mainlinePosition(ctx, e, index, positions)
		if err != nil {
			return nil, err
		}
		keys[id] = sortKey{rank: int64(position), ts: e.OriginServerTS(), id: id}
	}

	out := make([]*event.Event, 0, len(events))
	for _, e := range events {
		out = append(out, e)
	}
	sortByKey(out, keys)
	return out, nil
}

// noMainlinePosition is the spec's "infinity": an event with no power
// levels event behind it at all. It sorts before everything else, because
// it depends on a point in the room older than any power change.
const noMainlinePosition = math.MaxInt

// mainlinePosition returns how far back in the mainline an event depends.
//
// The event itself is not consulted, only the power levels events it
// reaches through its auth events — so a power levels event's own position
// is its predecessor's. That reading never comes up in practice: every
// power levels event is a power event, and power events are all settled in
// the first pass, so none of them reaches this sort.
//
// Each event cites at most one power levels event, so the walk follows a
// single path and every event along it shares the answer.
func (r *resolver) mainlinePosition(ctx context.Context, e *event.Event, index map[string]int, memo map[string]int) (int, error) {
	var walked []string
	current := e
	for {
		id := current.EventID()
		if position, ok := memo[id]; ok {
			return record(walked, position, memo), nil
		}
		next, err := r.powerLevelsInAuthEvents(ctx, current)
		if err != nil {
			return 0, err
		}
		if next == nil {
			return record(append(walked, id), noMainlinePosition, memo), nil
		}
		if position, ok := index[next.EventID()]; ok {
			return record(append(walked, id), position, memo), nil
		}
		walked = append(walked, id)
		current = next
	}
}

// record caches one position for every event on the path that led to it.
func record(walked []string, position int, memo map[string]int) int {
	for _, id := range walked {
		memo[id] = position
	}
	return position
}

// powerLevelsInAuthEvents returns the m.room.power_levels event an event
// cites, if it cites one.
func (r *resolver) powerLevelsInAuthEvents(ctx context.Context, e *event.Event) (*event.Event, error) {
	for _, id := range e.AuthEventIDs() {
		cited, err := r.get(ctx, id)
		if err != nil {
			return nil, err
		}
		if cited == nil {
			continue
		}
		if sk, isState := cited.StateKey(); isState && sk == "" && cited.Type() == event.TypePowerLevels {
			return cited, nil
		}
	}
	return nil, nil
}

// sortKey is what both orderings compare.
//
// rank holds the sender's power level for the reverse topological power
// ordering and the mainline position for the mainline ordering. The two
// orderings are otherwise the same comparison, and in both of them a
// higher rank sorts first: more power goes first, and so does a position
// further back in the mainline.
type sortKey struct {
	rank int64
	ts   int64
	id   string
}

func compareKeys(a, b sortKey) int {
	switch {
	case a.rank > b.rank:
		return -1
	case a.rank < b.rank:
		return 1
	case a.ts < b.ts:
		return -1
	case a.ts > b.ts:
		return 1
	default:
		// Event IDs are unique, so the order is total: there is never a
		// tie left for a server to break its own way.
		return strings.Compare(a.id, b.id)
	}
}

func sortByKey(events []*event.Event, keys map[string]sortKey) {
	slices.SortFunc(events, func(a, b *event.Event) int {
		return compareKeys(keys[a.EventID()], keys[b.EventID()])
	})
}

// sortHeap is the priority queue Kahn's algorithm picks from.
type sortHeap struct {
	ids  []string
	keys map[string]sortKey
}

func (h *sortHeap) Len() int           { return len(h.ids) }
func (h *sortHeap) Less(i, j int) bool { return compareKeys(h.keys[h.ids[i]], h.keys[h.ids[j]]) < 0 }
func (h *sortHeap) Swap(i, j int)      { h.ids[i], h.ids[j] = h.ids[j], h.ids[i] }
func (h *sortHeap) Push(x any)         { h.ids = append(h.ids, x.(string)) }

func (h *sortHeap) Pop() any {
	last := len(h.ids) - 1
	out := h.ids[last]
	h.ids = h.ids[:last]
	return out
}