package roomserver

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"sync"

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

// The intake of events from other servers.
//
// An event created here is built on the room's current state and can only
// fail the authorisation rules. An event from somewhere else arrives with
// none of that settled: it was built on whatever its sender knew, its auth
// events and its prev_events may be events this server has never seen,
// and it may be perfectly valid where it was made and wrong for the room as
// this server knows it. The spec's procedure for it, which this file
// follows step by step, is:
//
//  1. It is well-formed, signed by its sender's server, and its content
//     matches its hash — or it is replaced by its redacted form.
//  2. It passes the authorisation rules against its own auth_events,
//     which are fetched if they are missing.
//  3. It passes them against the room's state *before* it — the state
//     after its prev_events, resolved when there are several. A prev_event
//     this server does not have is asked for with /get_missing_events and,
//     if the gap is too wide, replaced by asking for the state at it.
//  4. It passes them against the room's *current* state. One that fails
//     only here is soft-failed: kept, so that other events can cite it, but
//     never shown to a client and never a forward extremity.
//
// Failing 1 to 3 is a rejection, and a rejected event is not stored.

// Bounds on how much one event can make this server fetch.
const (
	// missingEventsLimit is how many events /get_missing_events is asked
	// for. A gap wider than this is closed by fetching the state instead.
	missingEventsLimit = 10

	// stateFetchByIDLimit is how many missing state events are fetched one
	// at a time with /event before it is cheaper to ask for the whole
	// state with /state.
	stateFetchByIDLimit = 50

	// fetchConcurrency bounds the /event requests in flight at once.
	fetchConcurrency = 8
)

// ReceivePDU takes one PDU from a /send transaction. It returns the event
// ID — which a transaction's response is keyed by — and the reason the
// event was refused, if it was.
func (r *API) ReceivePDU(ctx context.Context, origin string, raw json.RawMessage) (string, error) {
	if err := r.requireFederation("receiving events from another server"); err != nil {
		return "", err
	}
	roomID, err := roomIDOfRaw(raw)
	if err != nil {
		return "", err
	}
	room, rules, err := r.roomAndRules(ctx, roomID)
	if errors.Is(err, ErrNoSuchRoom) {
		return "", fmt.Errorf("%w: %s", ErrNotInRoom, roomID)
	} else if err != nil {
		return "", err
	}
	e, err := parseRemote(raw, rules)
	if err != nil {
		return "", err
	}
	if in, err := r.participating(ctx, roomID); err != nil {
		return e.EventID(), err
	} else if !in {
		return e.EventID(), fmt.Errorf("%w: this server is not in %s", ErrNotInRoom, roomID)
	}
	if err := r.requireAllowed(ctx, roomID, origin); err != nil {
		return e.EventID(), err
	}
	return e.EventID(), r.receive(ctx, incoming{origin: origin, room: room, rules: rules, event: e, fetchMissing: true})
}

// incoming is one event on its way in.
type incoming struct {
	origin string
	room   storage.Room
	rules  roomversion.Rules
	event  *event.Event

	// fetchMissing allows /get_missing_events. It is off for the events
	// that call itself returns, so that a gap is closed one level deep and
	// no further.
	fetchMissing bool

	// relay marks an event handed over by send_join or send_leave; see
	// Notification.Relay.
	relay bool

	// batch is the part of a backfill batch already accepted, which is
	// not stored until the whole batch is (see backfill.go). The events in
	// it are this event's likeliest auth and prev events, so they are
	// looked for here before the store and before asking anybody.
	batch map[string]batchEvent
}

// receive runs the intake procedure.
func (r *API) receive(ctx context.Context, in incoming) error {
	// An event already held, as an outlier or in the timeline, has been
	// through this before. An outlier is left an outlier even when it
	// arrives again as an ordinary event: it is one because this server
	// used it as state without having its history, and making it a forward
	// extremity now would put the room's head on an event whose children
	// may already be here.
	if _, err := r.store.Event(ctx, in.room.ID, in.event.EventID()); err == nil {
		return nil
	} else if !isNotFound(err) {
		return err
	}

	e, err := r.verifyIncoming(ctx, in.event)
	if err != nil {
		return err
	}
	in.event = e

	// Step 2: the event against its own auth events.
	authEvents, err := r.authEventsOf(ctx, in)
	if err != nil {
		return err
	}
	create, err := r.createEvent(ctx, in.room, in.rules)
	if err != nil {
		return err
	}
	if err := auth.CheckAuthEvents(e, authEvents); err != nil {
		return err
	}
	if err := auth.Check(ctx, e, withCreate(auth.NewState(authEvents), create), r.verifier()); err != nil {
		return err
	}

	// Step 3: the event against the state before it.
	stateBefore, err := r.stateBeforeIncoming(ctx, in)
	if err != nil {
		return err
	}
	if err := auth.Check(ctx, e, withCreate(stateBefore, create), r.verifier()); err != nil {
		return fmt.Errorf("%w (against the state before it)", err)
	}

	return r.persistIncoming(ctx, in, stateBefore)
}

// withCreate adds the create event to a state that lacks it, which from
// room version 12 is every state built from auth_events.
func withCreate(st auth.State, create *event.Event) auth.State {
	if st.Create() == nil && create != nil {
		st[auth.StateKey{Type: event.TypeCreate}] = create
	}
	return st
}

// createEvent reads a room's create event.
func (r *API) createEvent(ctx context.Context, room storage.Room, rules roomversion.Rules) (*event.Event, error) {
	stored, err := r.store.CurrentStateEvent(ctx, room.ID, event.TypeCreate, "")
	if err != nil {
		if isNotFound(err) {
			return nil, fmt.Errorf("%w: this server does not hold the state of %s", ErrNotInRoom, room.ID)
		}
		return nil, err
	}
	return parseEvent(stored, rules)
}

// persistIncoming is step 4 and the write: under the room lock, the event
// is checked against the current state and either soft-failed or added to
// the room's head.
func (r *API) persistIncoming(ctx context.Context, in incoming, stateBefore auth.State) error {
	e := in.event
	mu := r.lock(in.room.ID)
	mu.Lock()
	defer mu.Unlock()

	// Another transaction may have delivered the same event while this one
	// was fetching what it needed.
	if _, err := r.store.Event(ctx, in.room.ID, e.EventID()); err == nil {
		return nil
	} else if !isNotFound(err) {
		return err
	}

	current, extremities, _, err := r.roomState(ctx, in.room, in.rules)
	if err != nil {
		return err
	}

	// Step 4: soft failure.
	if err := auth.Check(ctx, e, current, r.verifier()); err != nil {
		r.log.Info("soft-failed an event from another server",
			"room_id", in.room.ID, "event_id", e.EventID(), "type", e.Type(),
			"sender", e.Sender(), "origin", in.origin, "reason", err)
		return r.store.WriteOutliers(ctx, in.room.ID, []storage.StoredEvent{storedOf(in.room.ID, e)})
	}

	// The new forward extremities are the old ones this event does not
	// descend from, plus the event.
	prevs := e.PrevEventIDs()
	var newExtremities []string
	for _, x := range extremities {
		if !slices.Contains(prevs, x.EventID()) {
			newExtremities = append(newExtremities, x.EventID())
		}
	}
	newExtremities = append(newExtremities, e.EventID())

	req := persistRequest{
		Room:        in.room,
		Rules:       in.rules,
		Event:       e,
		StateBefore: stateBefore,
		Extremities: newExtremities,
		Origin:      in.origin,
		Relay:       in.relay,
	}
	if e.Type() == event.TypeRedaction {
		req.RedactsEventID = e.Redacts()
	}

	// The ordinary case: the event follows the room's single extremity, so
	// the state before it is the current state and the state after it is
	// one tuple more. Anything else — a fork, or an event built on an
	// older part of the room — means working the current state out again
	// from the room's new head.
	linear := len(extremities) == 1 && len(prevs) == 1 && prevs[0] == extremities[0].EventID()
	if !linear {
		resolved, err := r.stateAtHead(ctx, in.room, in.rules, newExtremities, e, stateBefore)
		if err != nil {
			return err
		}
		req.CurrentState = resolved
	}
	return r.persist(ctx, req)
}

// stateAtHead resolves the room's state across a set of forward
// extremities, one of which is the event being persisted.
func (r *API) stateAtHead(ctx context.Context, room storage.Room, rules roomversion.Rules, extremityIDs []string, e *event.Event, stateBefore auth.State) (auth.State, error) {
	var sets [][]*event.Event
	for _, id := range extremityIDs {
		if id == e.EventID() {
			set := stateres.Sorted(stateBefore)
			if e.IsState() {
				set = append(set, e)
			}
			sets = append(sets, set)
			continue
		}
		stored, err := r.store.Event(ctx, room.ID, id)
		if err != nil {
			return nil, err
		}
		x, err := parseEvent(stored, rules)
		if err != nil {
			return nil, err
		}
		set, err := r.stateAfter(ctx, room.ID, rules, x)
		if err != nil {
			return nil, err
		}
		sets = append(sets, set)
	}
	if len(sets) == 1 {
		return auth.NewState(sets[0]), nil
	}
	resolved, err := stateres.Resolve(ctx, rules, sets, eventProvider{store: r.store, roomID: room.ID, rules: rules})
	if err != nil {
		return nil, fmt.Errorf("resolving the state of %s: %w", room.ID, err)
	}
	return resolved, nil
}

// authEventsOf returns an incoming event's auth events, fetching the ones
// this server does not hold.
func (r *API) authEventsOf(ctx context.Context, in incoming) ([]*event.Event, error) {
	ids := in.event.AuthEventIDs()
	have, err := r.heldOrBatched(ctx, in, ids)
	if err != nil {
		return nil, err
	}
	if len(have) < len(ids) {
		if err := r.fetchAuthChain(ctx, in); err != nil {
			return nil, err
		}
		if have, err = r.heldOrBatched(ctx, in, ids); err != nil {
			return nil, err
		}
	}
	if len(have) < len(ids) {
		return nil, fmt.Errorf("%w: %s cites auth events that %s could not supply",
			ErrBadEvent, in.event.EventID(), in.origin)
	}

	// Keep the order the event cites them in; CheckAuthEvents does not
	// care, but a log line naming them does.
	byID := map[string]*event.Event{}
	for _, e := range have {
		byID[e.EventID()] = e
	}
	out := make([]*event.Event, 0, len(ids))
	for _, id := range ids {
		out = append(out, byID[id])
	}
	return out, nil
}

// heldOrBatched is eventsByID that also finds the events of the backfill
// batch an incoming event belongs to.
func (r *API) heldOrBatched(ctx context.Context, in incoming, ids []string) ([]*event.Event, error) {
	var out, rest []string
	for _, id := range ids {
		if _, ok := in.batch[id]; ok {
			out = append(out, id)
		} else {
			rest = append(rest, id)
		}
	}
	held, err := r.eventsByID(ctx, in.room.ID, in.rules, rest)
	if err != nil {
		return nil, err
	}
	for _, id := range out {
		held = append(held, in.batch[id].event)
	}
	return held, nil
}

// eventsByID reads and parses the events of a room this server holds,
// outliers included, leaving out the ones it does not.
func (r *API) eventsByID(ctx context.Context, roomID string, rules roomversion.Rules, ids []string) ([]*event.Event, error) {
	if len(ids) == 0 {
		return nil, nil
	}
	stored, err := r.store.EventsByID(ctx, roomID, ids)
	if err != nil {
		return nil, err
	}
	return parseEvents(stored, rules)
}

// fetchAuthChain asks the origin for an event's whole auth chain and
// stores what passes as outliers.
func (r *API) fetchAuthChain(ctx context.Context, in incoming) error {
	resp, err := r.remote.EventAuth(ctx, in.origin, in.room.ID, in.event.EventID())
	if err != nil {
		return remoteError(in.origin, err)
	}
	return r.acceptOutliers(ctx, in.room, in.rules, resp.AuthChain)
}

// acceptOutliers parses, verifies and authorises a batch of events from
// another server, and stores the ones that pass as outliers.
func (r *API) acceptOutliers(ctx context.Context, room storage.Room, rules roomversion.Rules, raws []json.RawMessage) error {
	var fresh []*event.Event
	referenced := map[string]bool{}
	ids := make([]string, 0, len(raws))
	for _, raw := range raws {
		e, err := parseRemote(raw, rules)
		if err != nil {
			r.log.Debug("ignoring an unparseable event", "room_id", room.ID, "error", err)
			continue
		}
		if !inRoom(e, room.ID) {
			continue
		}
		ids = append(ids, e.EventID())
		fresh = append(fresh, e)
		for _, id := range e.AuthEventIDs() {
			referenced[id] = true
		}
	}

	// What is already here does not need fetching, verifying or storing
	// again; it is only needed as the auth events of what is new.
	held, err := r.eventsByID(ctx, room.ID, rules, ids)
	if err != nil {
		return err
	}
	heldIDs := map[string]bool{}
	for _, e := range held {
		heldIDs[e.EventID()] = true
	}
	var toCheck []*event.Event
	for _, e := range fresh {
		if !heldIDs[e.EventID()] {
			toCheck = append(toCheck, e)
		}
	}
	if len(toCheck) == 0 {
		return nil
	}

	verified := r.verifyAll(ctx, toCheck)

	knownIDs := make([]string, 0, len(referenced))
	for id := range referenced {
		knownIDs = append(knownIDs, id)
	}
	knownEvents, err := r.eventsByID(ctx, room.ID, rules, knownIDs)
	if err != nil {
		return err
	}
	known := map[string]*event.Event{}
	for _, e := range knownEvents {
		known[e.EventID()] = e
	}
	if create, err := r.createEvent(ctx, room, rules); err == nil {
		known[create.EventID()] = create
	}

	passed, rejected := r.authCheckChain(ctx, verified, known)
	if len(rejected) > 0 {
		r.log.Info("rejected events in a fetched auth chain or state",
			"room_id", room.ID, "rejected", len(rejected))
	}
	rows := make([]storage.StoredEvent, 0, len(passed))
	for _, e := range verified {
		if _, ok := passed[e.EventID()]; ok {
			rows = append(rows, storedOf(room.ID, e))
		}
	}
	return r.store.WriteOutliers(ctx, room.ID, rows)
}

// verifyAll checks the signatures and hashes of a batch, in parallel, and
// returns the events that pass in their original order.
func (r *API) verifyAll(ctx context.Context, events []*event.Event) []*event.Event {
	out := make([]*event.Event, len(events))
	sem := make(chan struct{}, fetchConcurrency)
	var wg sync.WaitGroup
	for i, e := range events {
		wg.Add(1)
		sem <- struct{}{}
		go func() {
			defer wg.Done()
			defer func() { <-sem }()
			verified, err := r.verifyIncoming(ctx, e)
			if err != nil {
				r.log.Debug("dropping an event that does not verify", "event_id", e.EventID(), "error", err)
				return
			}
			out[i] = verified
		}()
	}
	wg.Wait()
	return slices.DeleteFunc(out, func(e *event.Event) bool { return e == nil })
}

// stateBeforeIncoming works out the room state an incoming event was built
// on: the state after each of its prev_events, resolved together.
func (r *API) stateBeforeIncoming(ctx context.Context, in incoming) (auth.State, error) {
	prevs := in.event.PrevEventIDs()
	if len(prevs) == 0 {
		return nil, fmt.Errorf("%w: %s has no prev_events", ErrBadEvent, in.event.EventID())
	}

	// A prev_event earlier in the same backfill batch has its state after
	// it in hand already.
	var sets [][]*event.Event
	var stored []string
	for _, id := range prevs {
		if b, ok := in.batch[id]; ok {
			sets = append(sets, b.stateAfter)
		} else {
			stored = append(stored, id)
		}
	}

	known, missing, err := r.splitPrevs(ctx, in.room, in.rules, stored)
	if err != nil {
		return nil, err
	}
	if len(missing) > 0 && in.fetchMissing {
		r.fetchMissingEvents(ctx, in)
		if known, missing, err = r.splitPrevs(ctx, in.room, in.rules, stored); err != nil {
			return nil, err
		}
	}

	for _, prev := range known {
		set, err := r.stateAfter(ctx, in.room.ID, in.rules, prev)
		if err != nil {
			return nil, err
		}
		sets = append(sets, set)
	}
	for _, prevID := range missing {
		set, err := r.stateAfterRemote(ctx, in, prevID)
		if err != nil {
			return nil, err
		}
		sets = append(sets, set)
	}

	if len(sets) == 1 {
		return auth.NewState(sets[0]), nil
	}
	resolved, err := stateres.Resolve(ctx, in.rules, sets,
		eventProvider{store: r.store, roomID: in.room.ID, rules: in.rules})
	if err != nil {
		return nil, fmt.Errorf("resolving the state before %s: %w", in.event.EventID(), err)
	}
	return resolved, nil
}

// splitPrevs sorts prev_events into the ones in this server's timeline,
// which have a state snapshot, and the ones it lacks. An outlier counts as
// lacking: it is held, but the state at it is not.
func (r *API) splitPrevs(ctx context.Context, room storage.Room, rules roomversion.Rules, prevs []string) ([]*event.Event, []string, error) {
	stored, err := r.store.EventsByID(ctx, room.ID, prevs)
	if err != nil {
		return nil, nil, err
	}
	have := map[string]storage.StoredEvent{}
	for _, s := range stored {
		if !s.Outlier {
			have[s.EventID] = s
		}
	}
	var known []*event.Event
	var missing []string
	for _, id := range prevs {
		s, ok := have[id]
		if !ok {
			missing = append(missing, id)
			continue
		}
		e, err := parseEvent(s, rules)
		if err != nil {
			return nil, nil, err
		}
		known = append(known, e)
	}
	return known, missing, nil
}

// fetchMissingEvents asks the origin for the events between this server's
// head of the room and the incoming event, and takes each of them in,
// oldest first. Failures are logged, not returned: whatever this closes of
// the gap helps, and what it leaves open is closed by fetching state.
func (r *API) fetchMissingEvents(ctx context.Context, in incoming) {
	stored, err := r.store.Extremities(ctx, in.room.ID)
	if err != nil {
		r.log.Warn("could not read extremities to fetch missing events", "room_id", in.room.ID, "error", err)
		return
	}
	earliest := make([]string, 0, len(stored))
	var minDepth int64
	for i, s := range stored {
		earliest = append(earliest, s.EventID)
		if i == 0 || s.Depth < minDepth {
			minDepth = s.Depth
		}
	}

	resp, err := r.remote.MissingEvents(ctx, in.origin, in.room.ID, federation.MissingEventsRequest{
		EarliestEvents: earliest,
		LatestEvents:   []string{in.event.EventID()},
		Limit:          missingEventsLimit,
		MinDepth:       minDepth,
	})
	if err != nil {
		r.log.Info("could not fetch missing events", "room_id", in.room.ID,
			"origin", in.origin, "event_id", in.event.EventID(), "error", err)
		return
	}

	var events []*event.Event
	for _, raw := range resp.Events {
		e, err := parseRemote(raw, in.rules)
		if err != nil || e.RoomID() != in.room.ID || e.EventID() == in.event.EventID() {
			continue
		}
		events = append(events, e)
	}
	eventsByDepth(events)
	for _, e := range events {
		err := r.receive(ctx, incoming{origin: in.origin, room: in.room, rules: in.rules, event: e})
		if err != nil {
			r.log.Info("a fetched missing event was not accepted",
				"room_id", in.room.ID, "event_id", e.EventID(), "error", err)
		}
	}
}

// stateAfterRemote asks the origin for the state at a prev_event this
// server does not have, fetches whatever of it is missing, and returns the
// state after that event.
//
// /state_ids answers with the state *before* the event it is asked about,
// so the event itself is fetched too and, if it is a state event, added on
// top. Everything fetched is authorised against its own auth events and
// stored as an outlier: it is real, but this server has none of the
// history that would put it in the timeline.
func (r *API) stateAfterRemote(ctx context.Context, in incoming, prevID string) ([]*event.Event, error) {
	ids, err := r.remote.StateIDs(ctx, in.origin, in.room.ID, prevID)
	if err != nil {
		return nil, remoteError(in.origin, err)
	}

	need := append(append(slices.Clone(ids.PDUIDs), ids.AuthChainIDs...), prevID)
	have, err := r.store.EventsByID(ctx, in.room.ID, need)
	if err != nil {
		return nil, err
	}
	haveIDs := map[string]bool{}
	for _, s := range have {
		haveIDs[s.EventID] = true
	}
	var missing []string
	for _, id := range need {
		if !haveIDs[id] && !slices.Contains(missing, id) {
			missing = append(missing, id)
		}
	}

	if len(missing) > 0 {
		var raws []json.RawMessage
		if len(missing) > stateFetchByIDLimit {
			state, err := r.remote.State(ctx, in.origin, in.room.ID, prevID)
			if err != nil {
				return nil, remoteError(in.origin, err)
			}
			raws = append(append(raws, state.AuthChain...), state.PDUs...)
			if !slices.Contains(missing, prevID) {
				missing = nil
			} else {
				missing = []string{prevID}
			}
		}
		raws = append(raws, r.fetchEvents(ctx, in.origin, missing)...)
		if err := r.acceptOutliers(ctx, in.room, in.rules, raws); err != nil {
			return nil, err
		}
	}

	state, err := r.eventsByID(ctx, in.room.ID, in.rules, ids.PDUIDs)
	if err != nil {
		return nil, err
	}
	prev, err := r.eventsByID(ctx, in.room.ID, in.rules, []string{prevID})
	if err != nil {
		return nil, err
	}
	if len(prev) == 1 && prev[0].IsState() {
		state = append(state, prev[0])
	}
	return state, nil
}

// fetchEvents fetches events one by one from a server, in parallel,
// returning the ones it supplied.
func (r *API) fetchEvents(ctx context.Context, server string, ids []string) []json.RawMessage {
	results := make([]json.RawMessage, len(ids))
	sem := make(chan struct{}, fetchConcurrency)
	var wg sync.WaitGroup
	for i, id := range ids {
		wg.Add(1)
		sem <- struct{}{}
		go func() {
			defer wg.Done()
			defer func() { <-sem }()
			txn, err := r.remote.Event(ctx, server, id)
			if err != nil || len(txn.PDUs) == 0 {
				r.log.Debug("could not fetch an event", "server", server, "event_id", id, "error", err)
				return
			}
			results[i] = txn.PDUs[0]
		}()
	}
	wg.Wait()
	return slices.DeleteFunc(results, func(raw json.RawMessage) bool { return raw == nil })
}