// Package stateres implements Matrix state resolution, version 2.
//
// Every event in a room names the events it follows, so a room is a graph
// rather than a line, and two servers can quite legitimately extend it at
// the same time. State resolution is the function that takes the differing
// views of the room's state at those points and produces one answer that
// every server computes identically. It is the reason two people who were
// both made moderator in the same second do not end up with the room in
// two different shapes.
//
// The algorithm is in the room version 2 specification, and room version
// 12 amends it. In outline:
//
//  1. State that every view agrees on is set aside as unconflicted.
//  2. The events in dispute, plus the events that authorise them, are
//     sorted so that whoever had the most power goes first.
//  3. They are replayed through the authorisation rules; anything the
//     rules refuse is dropped.
//  4. What is left is sorted again, this time by how far back in the
//     room's history of power changes it hangs, and replayed.
//  5. The unconflicted state is put back over the top.
//
// Sorting by power first is what makes the result resistant to a
// misbehaving server: an event that grants power can only win if the event
// that granted *it* power already won.
//
// The algorithm is deliberately total — it always produces a state — so
// nothing here rejects an event. Events the rules refuse are simply not
// applied.
package stateres

import (
	"context"
	"fmt"
	"slices"
	"strings"

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

// EventProvider supplies the events state resolution needs beyond the
// state sets themselves: the auth chains behind them, and the chain of
// power levels events used for the second sort.
//
// Events this server does not have, and events it rejected, must be left
// out of the result rather than reported as an error. That is not
// leniency: the spec requires rejected auth events to take no part in
// resolution, and leaving them out here is how a caller says so.
type EventProvider interface {
	EventsByID(ctx context.Context, ids []string) ([]*event.Event, error)
}

// Resolve resolves a set of room states into one.
//
// stateSets are the differing views of the room's state — in the usual
// case, the state after each of an event's prev_events. The result is the
// state they resolve to.
//
// Resolving fewer than two states is not an error and does no work: there
// is nothing to disagree about.
func Resolve(ctx context.Context, rules roomversion.Rules, stateSets [][]*event.Event, provider EventProvider) (auth.State, error) {
	switch len(stateSets) {
	case 0:
		return auth.State{}, nil
	case 1:
		return auth.NewState(stateSets[0]), nil
	}

	r := &resolver{
		rules:    rules,
		provider: provider,
		events:   map[string]*event.Event{},
		missing:  map[string]bool{},
		powerOf:  map[string]int64{},
	}
	maps := make([]auth.State, len(stateSets))
	for i, set := range stateSets {
		maps[i] = auth.NewState(set)
		for _, e := range set {
			r.events[e.EventID()] = e
		}
	}
	return r.resolve(ctx, maps)
}

// Sorted returns a state map as a slice in a deterministic order, for
// storing or comparing.
func Sorted(state auth.State) []*event.Event {
	keys := make([]auth.StateKey, 0, len(state))
	for k := range state {
		keys = append(keys, k)
	}
	slices.SortFunc(keys, func(a, b auth.StateKey) int {
		if a.Type != b.Type {
			return strings.Compare(a.Type, b.Type)
		}
		return strings.Compare(a.StateKey, b.StateKey)
	})
	out := make([]*event.Event, 0, len(keys))
	for _, k := range keys {
		out = append(out, state[k])
	}
	return out
}

// resolver holds the caches one resolution needs. It is single-use.
type resolver struct {
	rules    roomversion.Rules
	provider EventProvider

	// events is every event fetched so far, and missing records the IDs
	// that resolved to nothing, so that a dangling reference is looked up
	// once rather than on every walk that reaches it.
	events  map[string]*event.Event
	missing map[string]bool

	// powerOf caches the sender's power level per event, which the first
	// sort asks for repeatedly.
	powerOf map[string]int64
}

func (r *resolver) resolve(ctx context.Context, stateSets []auth.State) (auth.State, error) {
	unconflicted, conflicted := separate(stateSets)

	// The full conflicted set is wider than the events actually in
	// dispute: it also takes in the auth difference, the events that only
	// *some* of the views rely on. Without them, a view could smuggle in a
	// power grant by keeping it out of the state it declares.
	full := map[string]*event.Event{}
	for id, e := range conflicted {
		full[id] = e
	}
	authDiff, err := r.authDifference(ctx, stateSets)
	if err != nil {
		return nil, err
	}
	for id, e := range authDiff {
		full[id] = e
	}

	// Room version 12 adds the events that lie between conflicted events
	// in the auth graph, which the auth difference can miss. Both this and
	// the empty starting state below exist to make state resets rarer.
	if r.rules.StateRes == roomversion.StateResV2Subgraph {
		subgraph, err := r.conflictedSubgraph(ctx, conflicted)
		if err != nil {
			return nil, err
		}
		for id, e := range subgraph {
			full[id] = e
		}
	}

	// Step 1: the power events in the full conflicted set, together with
	// whatever in their auth chains is also in dispute.
	power, err := r.powerEventsToSort(ctx, full)
	if err != nil {
		return nil, err
	}
	sortedPower, err := r.reverseTopologicalPowerSort(ctx, power)
	if err != nil {
		return nil, err
	}

	// Step 2. Room version 12 starts from an empty state rather than the
	// unconflicted one: an event then has to be authorised by the events
	// it actually cites, which stops an unconflicted event from propping
	// up a conflicted one that nothing else supports.
	start := unconflicted
	if r.rules.StateRes == roomversion.StateResV2Subgraph {
		start = auth.State{}
	}
	partial, err := r.iterativeAuthChecks(ctx, start, sortedPower)
	if err != nil {
		return nil, err
	}

	// Step 3: everything else, ordered by where it hangs off the room's
	// history of power levels events.
	rest := make(map[string]*event.Event, len(full))
	for id, e := range full {
		if _, done := power[id]; !done {
			rest[id] = e
		}
	}
	sortedRest, err := r.mainlineSort(ctx, rest, partial.PowerLevels())
	if err != nil {
		return nil, err
	}

	// Step 4.
	resolved, err := r.iterativeAuthChecks(ctx, partial, sortedRest)
	if err != nil {
		return nil, err
	}

	// Step 5: the unconflicted state was never in doubt, so it wins
	// outright — including over anything the replay above put in its place.
	for k, e := range unconflicted {
		resolved[k] = e
	}
	return resolved, nil
}

// separate splits the state sets into the state they all agree on and the
// state they do not.
//
// A key missing from one view counts as a disagreement, not as a blank:
// one server thinking there is no topic and another thinking there is one
// is exactly the kind of conflict this algorithm exists to settle.
func separate(stateSets []auth.State) (auth.State, map[string]*event.Event) {
	unconflicted := auth.State{}
	conflicted := map[string]*event.Event{}

	keys := map[auth.StateKey]bool{}
	for _, set := range stateSets {
		for k := range set {
			keys[k] = true
		}
	}

	for k := range keys {
		var distinct []*event.Event
		seen := map[string]bool{}
		absent := false
		for _, set := range stateSets {
			e, ok := set[k]
			if !ok {
				absent = true
				continue
			}
			if !seen[e.EventID()] {
				seen[e.EventID()] = true
				distinct = append(distinct, e)
			}
		}
		if len(distinct) == 1 && !absent {
			unconflicted[k] = distinct[0]
			continue
		}
		for _, e := range distinct {
			conflicted[e.EventID()] = e
		}
	}
	return unconflicted, conflicted
}

// authDifference returns the events that are in some views' auth chains but
// not in all of them.
func (r *resolver) authDifference(ctx context.Context, stateSets []auth.State) (map[string]*event.Event, error) {
	chains := make([]map[string]bool, len(stateSets))
	for i, set := range stateSets {
		ids := make([]string, 0, len(set))
		for _, e := range set {
			ids = append(ids, e.EventID())
		}
		chain, err := r.authChain(ctx, ids)
		if err != nil {
			return nil, err
		}
		chains[i] = chain
	}

	diff := map[string]*event.Event{}
	for i, chain := range chains {
		for id := range chain {
			inAll := true
			for j, other := range chains {
				if i != j && !other[id] {
					inAll = false
					break
				}
			}
			if inAll {
				continue
			}
			if e := r.events[id]; e != nil {
				diff[id] = e
			}
		}
	}
	return diff, nil
}

// authChain returns the transitive closure of the auth_events links from
// the given events, not including the events themselves.
func (r *resolver) authChain(ctx context.Context, ids []string) (map[string]bool, error) {
	chain := map[string]bool{}
	queue := make([]string, 0, len(ids))
	for _, id := range ids {
		e, err := r.get(ctx, id)
		if err != nil {
			return nil, err
		}
		if e != nil {
			queue = append(queue, e.AuthEventIDs()...)
		}
	}

	for len(queue) > 0 {
		id := queue[len(queue)-1]
		queue = queue[:len(queue)-1]
		if chain[id] {
			continue
		}
		e, err := r.get(ctx, id)
		if err != nil {
			return nil, err
		}
		if e == nil {
			// A reference we cannot follow. It is not an error: the event
			// may have been rejected, or we may simply never have seen it.
			continue
		}
		chain[id] = true
		queue = append(queue, e.AuthEventIDs()...)
	}
	return chain, nil
}

// conflictedSubgraph returns the events that lie on an auth_events path
// between two conflicted events, which room version 12 adds to the full
// conflicted set.
//
// It is computed in two passes: walk forwards from the conflicted events
// to find everything they reach, recording the edges in reverse, then walk
// those edges backwards from the conflicted events. A node the second walk
// reaches is one that both descends from a conflicted event and leads to
// one, which is what being on a path between two of them means.
func (r *resolver) conflictedSubgraph(ctx context.Context, conflicted map[string]*event.Event) (map[string]*event.Event, error) {
	parents := map[string][]string{}
	seen := map[string]bool{}
	queue := make([]string, 0, len(conflicted))
	for id := range conflicted {
		queue = append(queue, id)
		seen[id] = true
	}

	for len(queue) > 0 {
		id := queue[len(queue)-1]
		queue = queue[:len(queue)-1]
		e, err := r.get(ctx, id)
		if err != nil {
			return nil, err
		}
		if e == nil {
			continue
		}
		for _, child := range e.AuthEventIDs() {
			parents[child] = append(parents[child], id)
			if !seen[child] {
				seen[child] = true
				queue = append(queue, child)
			}
		}
	}

	subgraph := map[string]*event.Event{}
	visited := map[string]bool{}
	queue = queue[:0]
	for id := range conflicted {
		queue = append(queue, id)
		visited[id] = true
	}
	for len(queue) > 0 {
		id := queue[len(queue)-1]
		queue = queue[:len(queue)-1]
		if e := r.events[id]; e != nil {
			subgraph[id] = e
		}
		for _, parent := range parents[id] {
			if !visited[parent] {
				visited[parent] = true
				queue = append(queue, parent)
			}
		}
	}
	return subgraph, nil
}

// powerEventsToSort returns the events the first pass sorts and replays:
// the power events in the full conflicted set, plus anything in their auth
// chains that is also in the full conflicted set.
func (r *resolver) powerEventsToSort(ctx context.Context, full map[string]*event.Event) (map[string]*event.Event, error) {
	out := map[string]*event.Event{}
	for id, e := range full {
		if !isPowerEvent(e) {
			continue
		}
		out[id] = e
		chain, err := r.authChain(ctx, []string{id})
		if err != nil {
			return nil, err
		}
		for chainID := range chain {
			if e, ok := full[chainID]; ok {
				out[chainID] = e
			}
		}
	}
	return out, nil
}

// isPowerEvent reports whether an event could take away somebody's ability
// to do something in the room, which is what makes it worth sorting first.
func isPowerEvent(e *event.Event) bool {
	switch e.Type() {
	case event.TypePowerLevels, event.TypeJoinRules:
		_, isState := e.StateKey()
		return isState
	case event.TypeMember:
		sk, ok := e.StateKey()
		if !ok || sk == e.Sender() {
			// Leaving or being banned under your own name takes nothing
			// from anybody else.
			return false
		}
		membership, err := e.Membership()
		if err != nil {
			return false
		}
		return membership == event.MembershipLeave || membership == event.MembershipBan
	default:
		return false
	}
}

// iterativeAuthChecks replays events through the authorisation rules,
// applying the ones the rules allow and ignoring the rest.
//
// The state each event is judged against is its own auth_events, overlaid
// with whatever the replay has settled so far for the keys the rules will
// read. So an event is authorised by the room as it is being rebuilt,
// falling back to the events it cited for anything not yet decided.
func (r *resolver) iterativeAuthChecks(ctx context.Context, base auth.State, events []*event.Event) (auth.State, error) {
	state := make(auth.State, len(base))
	for k, e := range base {
		state[k] = e
	}

	for _, e := range events {
		checkState := auth.State{}
		for _, id := range e.AuthEventIDs() {
			cited, err := r.get(ctx, id)
			if err != nil {
				return nil, err
			}
			if cited == nil {
				// Either we do not have it or we rejected it. A rejected
				// auth event must not be used, and this is where that
				// happens.
				continue
			}
			if sk, ok := cited.StateKey(); ok {
				checkState[auth.StateKey{Type: cited.Type(), StateKey: sk}] = cited
			}
		}
		keys, err := auth.AuthEventKeys(e)
		if err != nil {
			// A malformed event cannot be authorised, which in this
			// algorithm means it is skipped rather than fatal.
			continue
		}
		for _, k := range keys {
			if current, ok := state[k]; ok {
				checkState[k] = current
			}
		}

		// From room version 12 the create event is in neither the event's
		// auth_events nor the keys above, but the rules still need it: it
		// says who the creators are, and the room ID is checked against
		// it. Every other version has it from auth_events already.
		if checkState.Create() == nil {
			create, err := r.createEvent(ctx, e)
			if err != nil {
				return nil, err
			}
			if create != nil {
				checkState[auth.StateKey{Type: event.TypeCreate}] = create
			}
		}

		if err := auth.Check(ctx, e, checkState, intakeVerified{}); err != nil {
			continue
		}
		if sk, ok := e.StateKey(); ok {
			state[auth.StateKey{Type: e.Type(), StateKey: sk}] = e
		}
	}
	return state, nil
}

// intakeVerified stands in for the signature check the restricted join
// rule asks for.
//
// Every event that reaches state resolution was accepted into the room
// first, and that is where its signatures — including the one from the
// server authorising a restricted join — were checked. Checking it again
// here would mean fetching another server's keys in the middle of an
// algorithm that has to be deterministic, run identically on every server
// and return without doing any I/O it can avoid. Synapse draws the line in
// the same place.
//
// The consequence is that intake must not be skipped: an event inserted
// into storage without a signature check would be authorised here as
// though it had one.
type intakeVerified struct{}

func (intakeVerified) VerifyEventSignature(context.Context, string, *event.Event) error { return nil }

// createEvent returns the room's m.room.create event, as seen from one
// event in it.
//
// Up to room version 11 every event cites it in auth_events. From version
// 12 it is deliberately left out, because the room ID is its event ID with
// a different sigil — so it is looked up from there instead, which is the
// same guarantee by a shorter route.
func (r *resolver) createEvent(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.TypeCreate {
			return cited, nil
		}
	}
	if !r.rules.RoomIDIsCreateEventID {
		return nil, nil
	}
	roomID := e.RoomID()
	if len(roomID) < 2 || roomID[0] != '!' {
		return nil, nil
	}
	return r.get(ctx, "$"+roomID[1:])
}

// get returns an event by ID, from the cache or from the provider.
func (r *resolver) get(ctx context.Context, id string) (*event.Event, error) {
	if e, ok := r.events[id]; ok {
		return e, nil
	}
	if r.missing[id] {
		return nil, nil
	}
	fetched, err := r.provider.EventsByID(ctx, []string{id})
	if err != nil {
		return nil, fmt.Errorf("state resolution: fetch %s: %w", id, err)
	}
	for _, e := range fetched {
		r.events[e.EventID()] = e
	}
	if e, ok := r.events[id]; ok {
		return e, nil
	}
	r.missing[id] = true
	return nil, nil
}

// senderPowerLevel returns the power level of an event's sender, judged
// from the event's own auth_events rather than from the room as it is now.
//
// Judging it from the event's own view is the point: it is what the sender
// could do at the time, which is the only thing every server can agree on
// after the fact.
func (r *resolver) senderPowerLevel(ctx context.Context, e *event.Event) (int64, error) {
	if level, ok := r.powerOf[e.EventID()]; ok {
		return level, nil
	}

	var plEvent *event.Event
	for _, id := range e.AuthEventIDs() {
		cited, err := r.get(ctx, id)
		if err != nil {
			return 0, err
		}
		if cited == nil {
			continue
		}
		if sk, isState := cited.StateKey(); isState && sk == "" && cited.Type() == event.TypePowerLevels {
			plEvent = cited
			break
		}
	}
	createEvent, err := r.createEvent(ctx, e)
	if err != nil {
		return 0, err
	}

	level := int64(0)
	if createEvent != nil || plEvent != nil {
		pl, err := auth.ResolvePowerLevels(plEvent, createEvent, r.rules)
		if err != nil {
			// A power levels event we cannot read gives its sender
			// nothing, which keeps the sort total.
			level = 0
		} else {
			level = pl.UserLevel(e.Sender())
		}
	}
	r.powerOf[e.EventID()] = level
	return level, nil
}