package roomserver

import (
	"cmp"
	"context"
	"encoding/json"
	"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/id"
)

// Spaces: a space is a room whose create event says "type": "m.space",
// and whose m.space.child state events name the rooms in it. Nothing about
// that is enforced by the authorisation rules — a child event is a pointer,
// and the room it points at need not know — so what the server provides is
// the walk: /hierarchy reads a space's tree breadth first, summarising each
// room the user is allowed to know about, and asks the space's other
// servers about the rooms this one is not in.
//
// The same walk, one level deep and judged for a server rather than a user,
// answers the federation endpoint other servers walk Warp's spaces with.

// The space event types.
const (
	TypeSpaceChild  = "m.space.child"
	TypeSpaceParent = "m.space.parent"

	// RoomTypeSpace is m.room.create's content.type for a space.
	RoomTypeSpace = "m.space"
)

// Limits of a hierarchy walk.
const (
	// DefaultHierarchyLimit is the page size when the client names none,
	// and MaxHierarchyLimit caps the one it names.
	DefaultHierarchyLimit = 50
	MaxHierarchyLimit     = 100

	// MaxHierarchyDepth caps max_depth: a space tree deeper than this is
	// either a mistake or an attempt to make the server walk forever.
	MaxHierarchyDepth = 50

	// maxHierarchyRooms caps how many rooms one walk visits, whatever the
	// page size: paging re-walks from the root, so the cap is what bounds
	// the work of the last page of a large space.
	maxHierarchyRooms = 2000

	// maxOrderLength is the longest m.space.child order the spec lets
	// sort a child; a longer one is ignored, as if absent.
	maxOrderLength = 50
)

// SpaceChildEvent is an m.space.child event as a hierarchy carries it: the
// stripped state event plus its timestamp, which children are sorted by.
type SpaceChildEvent struct {
	Type           string          `json:"type"`
	StateKey       string          `json:"state_key"`
	Sender         string          `json:"sender"`
	Content        json.RawMessage `json:"content"`
	OriginServerTS int64           `json:"origin_server_ts"`
}

// SpaceRoom is one room of a hierarchy: the public face of the room, as
// the public room directory shows it, and the space's children when it has
// any.
type SpaceRoom struct {
	RoomID           string `json:"room_id"`
	Name             string `json:"name,omitempty"`
	Topic            string `json:"topic,omitempty"`
	AvatarURL        string `json:"avatar_url,omitempty"`
	CanonicalAlias   string `json:"canonical_alias,omitempty"`
	NumJoinedMembers int    `json:"num_joined_members"`
	WorldReadable    bool   `json:"world_readable"`
	GuestCanJoin     bool   `json:"guest_can_join"`
	JoinRule         string `json:"join_rule,omitempty"`
	RoomType         string `json:"room_type,omitempty"`
	RoomVersion      string `json:"room_version,omitempty"`
	Encryption       string `json:"encryption,omitempty"`

	// ChildrenState is the room's m.space.child events. It is always there
	// on the rooms a client is given, empty for a room that is not a
	// space, and left out of the children a federation response lists.
	ChildrenState []SpaceChildEvent `json:"children_state,omitempty"`

	// AllowedRoomIDs is the rooms whose members may join under a
	// restricted join rule. Only a federation response carries it: the
	// asking server needs it to decide what its own users may see.
	AllowedRoomIDs []string `json:"allowed_room_ids,omitempty"`
}

// HierarchyRequest is one page of a space's tree.
type HierarchyRequest struct {
	RoomID string

	// User is who is walking. Each room is shown only if they may know
	// about it.
	User string

	SuggestedOnly bool

	// MaxDepth is how far below the root to go; zero is the root alone.
	// Negative means the server's own limit.
	MaxDepth int

	// From is how many rooms of the walk earlier pages have already
	// returned.
	From int

	Limit int
}

// HierarchyPage is one page of a space's tree.
type HierarchyPage struct {
	Rooms []SpaceRoom

	// Next is the From of the next page, zero at the end.
	Next int
}

// FederationHierarchy is what GET /_matrix/federation/v1/hierarchy
// answers: the room, the children this server can describe, and the ones
// the asking server may not see.
type FederationHierarchy struct {
	Room                 SpaceRoom   `json:"room"`
	Children             []SpaceRoom `json:"children"`
	InaccessibleChildren []string    `json:"inaccessible_children"`
}

// spaceChild is one parsed m.space.child of a room.
type spaceChild struct {
	roomID    string
	via       []string
	order     string
	ts        int64
	suggested bool
}

// localSpaceRoom is a room this server is in, read for a hierarchy.
type localSpaceRoom struct {
	room     SpaceRoom
	children []spaceChild
	state    auth.State
}

// Hierarchy reads one page of a space's tree.
func (r *API) Hierarchy(ctx context.Context, req HierarchyRequest) (HierarchyPage, error) {
	maxDepth := req.MaxDepth
	if maxDepth < 0 || maxDepth > MaxHierarchyDepth {
		maxDepth = MaxHierarchyDepth
	}
	limit := req.Limit
	if limit <= 0 {
		limit = DefaultHierarchyLimit
	}
	limit = min(limit, MaxHierarchyLimit)

	type queued struct {
		roomID string
		depth  int
		via    []string
	}
	queue := []queued{{roomID: req.RoomID}}
	visited := map[string]bool{}
	remote := map[string]*FederationHierarchy{}

	var (
		out     HierarchyPage
		emitted int
	)
	for len(queue) > 0 && len(visited) < maxHierarchyRooms {
		next := queue[0]
		queue = queue[1:]
		if visited[next.roomID] {
			continue
		}
		visited[next.roomID] = true
		isRoot := next.roomID == req.RoomID

		room, children, accessible, err := r.hierarchyRoom(ctx, req.User, next.roomID, next.via, req.SuggestedOnly, remote)
		switch {
		case err != nil && isRoot:
			return HierarchyPage{}, err
		case err != nil:
			// A child this server cannot describe — not here, and no
			// server it names would say — is left out, as the spec has
			// it, rather than failing the walk.
			r.log.Debug("left a room out of a space hierarchy",
				"space", req.RoomID, "room_id", next.roomID, "error", err)
			continue
		case !accessible && isRoot:
			return HierarchyPage{}, fmt.Errorf("%w: %s may not see the space %s", ErrForbidden, req.User, req.RoomID)
		case !accessible:
			continue
		}

		if emitted >= req.From {
			if len(out.Rooms) == limit {
				out.Next = emitted
				return out, nil
			}
			out.Rooms = append(out.Rooms, room)
		}
		emitted++

		// Only a space's children are walked into: an ordinary room with
		// an m.space.child event is not a space, and the event in it is
		// an orphan.
		if next.depth >= maxDepth || (!isRoot && room.RoomType != RoomTypeSpace) {
			continue
		}
		for _, child := range children {
			if !visited[child.roomID] {
				queue = append(queue, queued{roomID: child.roomID, depth: next.depth + 1, via: child.via})
			}
		}
	}
	return out, nil
}

// hierarchyRoom describes one room of a walk, and reports whether the user
// may know about it.
func (r *API) hierarchyRoom(ctx context.Context, userID, roomID string, via []string, suggestedOnly bool, remote map[string]*FederationHierarchy) (SpaceRoom, []spaceChild, bool, error) {
	local, err := r.participating(ctx, roomID)
	if err != nil {
		return SpaceRoom{}, nil, false, err
	}
	if local {
		read, err := r.readSpaceRoom(ctx, roomID, suggestedOnly)
		if err != nil {
			return SpaceRoom{}, nil, false, err
		}
		accessible, err := r.userMayPeek(ctx, read, userID)
		if err != nil {
			return SpaceRoom{}, nil, false, err
		}
		read.room.AllowedRoomIDs = nil
		return read.room, read.children, accessible, nil
	}

	answer, err := r.remoteHierarchy(ctx, roomID, via, suggestedOnly, remote)
	if err != nil {
		return SpaceRoom{}, nil, false, err
	}
	room := answer.Room
	children := parseRemoteChildren(room.ChildrenState, suggestedOnly)
	if room.ChildrenState == nil {
		room.ChildrenState = []SpaceChildEvent{}
	}

	// The other server judged the room for this server as a whole; a
	// restricted room is shown to the user only if they are in one of the
	// rooms it lets in, which only this server knows.
	accessible := true
	if isRestricted(room.JoinRule) && !isPeekable(room.JoinRule, room.WorldReadable) {
		accessible = false
		for _, allowed := range room.AllowedRoomIDs {
			membership, err := r.store.Membership(ctx, allowed, userID)
			if err != nil {
				return SpaceRoom{}, nil, false, err
			}
			if membership.Membership == event.MembershipJoin {
				accessible = true
				break
			}
		}
	}
	room.AllowedRoomIDs = nil
	return room, children, accessible, nil
}

// remoteHierarchy asks the servers a child event names about a room,
// remembering the answer for the rest of the walk.
func (r *API) remoteHierarchy(ctx context.Context, roomID string, via []string, suggestedOnly bool, cache map[string]*FederationHierarchy) (*FederationHierarchy, error) {
	if cached, ok := cache[roomID]; ok {
		if cached == nil {
			return nil, fmt.Errorf("%w: %s", ErrNoSuchRoom, roomID)
		}
		return cached, nil
	}
	cache[roomID] = nil
	if !r.federating() || r.remote == nil {
		return nil, fmt.Errorf("%w: %s is not a room this server is in", ErrNoSuchRoom, roomID)
	}

	var lastErr error
	for _, server := range via {
		if server == r.serverName {
			continue
		}
		raw, err := r.remote.Hierarchy(ctx, server, roomID, suggestedOnly)
		if err != nil {
			lastErr = remoteError(server, err)
			continue
		}
		var answer FederationHierarchy
		if err := json.Unmarshal(raw, &answer); err != nil || answer.Room.RoomID != roomID {
			lastErr = fmt.Errorf("%w: %s sent an unreadable hierarchy for %s", ErrRemote, server, roomID)
			continue
		}
		cache[roomID] = &answer
		return &answer, nil
	}
	if lastErr == nil {
		lastErr = fmt.Errorf("%w: no server to ask about %s", ErrNoSuchRoom, roomID)
	}
	return nil, lastErr
}

// FederatedHierarchy answers another server walking a space of this one.
func (r *API) FederatedHierarchy(ctx context.Context, origin, roomID string, suggestedOnly bool) (FederationHierarchy, error) {
	local, err := r.participating(ctx, roomID)
	if err != nil {
		return FederationHierarchy{}, err
	}
	if !local {
		return FederationHierarchy{}, fmt.Errorf("%w: %s", ErrNoSuchRoom, roomID)
	}
	root, err := r.readSpaceRoom(ctx, roomID, suggestedOnly)
	if err != nil {
		return FederationHierarchy{}, err
	}
	accessible, err := r.serverMayPeek(ctx, root, origin)
	if err != nil {
		return FederationHierarchy{}, err
	}
	if !accessible {
		// Not found rather than forbidden: whether this server is in a
		// room the asker may not see is itself not theirs to know.
		return FederationHierarchy{}, fmt.Errorf("%w: %s", ErrNoSuchRoom, roomID)
	}

	out := FederationHierarchy{Room: root.room, Children: []SpaceRoom{}, InaccessibleChildren: []string{}}
	for _, child := range root.children {
		local, err := r.participating(ctx, child.roomID)
		if err != nil {
			return FederationHierarchy{}, err
		}
		if !local {
			continue
		}
		read, err := r.readSpaceRoom(ctx, child.roomID, suggestedOnly)
		if err != nil {
			return FederationHierarchy{}, err
		}
		ok, err := r.serverMayPeek(ctx, read, origin)
		if err != nil {
			return FederationHierarchy{}, err
		}
		if !ok {
			out.InaccessibleChildren = append(out.InaccessibleChildren, child.roomID)
			continue
		}
		read.room.ChildrenState = nil
		out.Children = append(out.Children, read.room)
	}
	return out, nil
}

// RoomSummaryRequest asks for the public face of one room.
type RoomSummaryRequest struct {
	RoomID string

	// User is who is asking, empty for a guest without a token.
	User string

	// Via are servers to ask when this one is not in the room.
	Via []string
}

// RoomSummaryResponse is GET /room_summary's answer: a hierarchy room
// without its children, plus where the user stands.
type RoomSummaryResponse struct {
	SpaceRoom
	Membership string `json:"membership,omitempty"`
}

// RoomSummary describes one room to a user who may not be in it, as a link
// preview or an invite needs it.
func (r *API) RoomSummary(ctx context.Context, req RoomSummaryRequest) (RoomSummaryResponse, error) {
	room, _, accessible, err := r.hierarchyRoom(ctx, req.User, req.RoomID, req.Via, false, map[string]*FederationHierarchy{})
	if err != nil {
		return RoomSummaryResponse{}, err
	}
	if !accessible {
		return RoomSummaryResponse{}, fmt.Errorf("%w: %s may not see %s", ErrForbidden, req.User, req.RoomID)
	}
	room.ChildrenState = nil
	out := RoomSummaryResponse{SpaceRoom: room}
	if req.User != "" {
		membership, err := r.store.Membership(ctx, req.RoomID, req.User)
		if err != nil && !isNotFound(err) {
			return RoomSummaryResponse{}, err
		}
		out.Membership = membership.Membership
		if out.Membership == "" {
			out.Membership = event.MembershipLeave
		}
	}
	return out, nil
}

// readSpaceRoom reads a room this server is in.
func (r *API) readSpaceRoom(ctx context.Context, roomID string, suggestedOnly bool) (localSpaceRoom, error) {
	room, rules, err := r.roomAndRules(ctx, roomID)
	if err != nil {
		return localSpaceRoom{}, err
	}
	summary, err := r.store.RoomSummary(ctx, roomID)
	if err != nil {
		return localSpaceRoom{}, err
	}
	st, err := r.currentState(ctx, room, rules)
	if err != nil {
		return localSpaceRoom{}, err
	}

	out := localSpaceRoom{
		state: st,
		room: SpaceRoom{
			RoomID:           roomID,
			Name:             summary.Name,
			Topic:            summary.Topic,
			AvatarURL:        summary.AvatarURL,
			CanonicalAlias:   summary.CanonicalAlias,
			NumJoinedMembers: summary.JoinedMembers,
			WorldReadable:    summary.HistoryVisibility == HistoryVisibilityWorldReadable,
			GuestCanJoin:     summary.GuestAccess == "can_join",
			JoinRule:         cmp.Or(st.JoinRule(), auth.JoinRuleInvite),
			RoomType:         summary.RoomType,
			RoomVersion:      room.Version,
			ChildrenState:    []SpaceChildEvent{},
		},
	}
	if enc := st.Get(TypeEncryption, ""); enc != nil {
		var content struct {
			Algorithm string `json:"algorithm"`
		}
		if json.Unmarshal(enc.Content(), &content) == nil {
			out.room.Encryption = content.Algorithm
		}
	}
	if isRestricted(out.room.JoinRule) {
		if out.room.AllowedRoomIDs, err = allowedRooms(st); err != nil {
			return localSpaceRoom{}, err
		}
	}

	for key, e := range st {
		if key.Type != TypeSpaceChild {
			continue
		}
		child, ok := parseSpaceChild(key.StateKey, e.Content(), e.OriginServerTS())
		if !ok || (suggestedOnly && !child.suggested) {
			continue
		}
		out.children = append(out.children, child)
		out.room.ChildrenState = append(out.room.ChildrenState, SpaceChildEvent{
			Type:           TypeSpaceChild,
			StateKey:       key.StateKey,
			Sender:         e.Sender(),
			Content:        e.Content(),
			OriginServerTS: e.OriginServerTS(),
		})
	}
	sortChildren(out.children)
	slices.SortFunc(out.room.ChildrenState, func(a, b SpaceChildEvent) int {
		return strings.Compare(a.StateKey, b.StateKey)
	})
	return out, nil
}

// userMayPeek reports whether a user may know about a room they are not
// necessarily in: the spec's accessibility rule for a hierarchy.
func (r *API) userMayPeek(ctx context.Context, read localSpaceRoom, userID string) (bool, error) {
	switch read.state.Membership(userID) {
	case event.MembershipJoin, event.MembershipInvite:
		return true, nil
	}
	if isPeekable(read.room.JoinRule, read.room.WorldReadable) {
		return true, nil
	}
	if isRestricted(read.room.JoinRule) {
		for _, allowed := range read.room.AllowedRoomIDs {
			membership, err := r.store.Membership(ctx, allowed, userID)
			if err != nil {
				return false, err
			}
			if membership.Membership == event.MembershipJoin {
				return true, nil
			}
		}
	}
	return false, nil
}

// serverMayPeek is userMayPeek for another server: anything a user of it
// could be shown. A restricted room is shown with its allow list, which
// the asking server checks against its own users.
func (r *API) serverMayPeek(ctx context.Context, read localSpaceRoom, origin string) (bool, error) {
	if isPeekable(read.room.JoinRule, read.room.WorldReadable) || isRestricted(read.room.JoinRule) {
		return true, nil
	}
	for key, e := range read.state {
		if key.Type != event.TypeMember {
			continue
		}
		server, err := id.ServerNameOf(key.StateKey)
		if err != nil || server != origin {
			continue
		}
		if membership, err := e.Membership(); err == nil &&
			(membership == event.MembershipJoin || membership == event.MembershipInvite) {
			return true, nil
		}
	}
	return false, nil
}

// isPeekable reports whether anybody at all may know about a room.
func isPeekable(joinRule string, worldReadable bool) bool {
	switch joinRule {
	case auth.JoinRulePublic, auth.JoinRuleKnock, auth.JoinRuleKnockRestricted:
		return true
	}
	return worldReadable
}

func isRestricted(joinRule string) bool {
	return joinRule == auth.JoinRuleRestricted || joinRule == auth.JoinRuleKnockRestricted
}

// parseSpaceChild reads one m.space.child. A child event without a
// non-empty `via` is how a space removes a child — the event cannot be
// deleted, only emptied — so it does not count.
func parseSpaceChild(stateKey string, content json.RawMessage, ts int64) (spaceChild, bool) {
	if !strings.HasPrefix(stateKey, "!") {
		return spaceChild{}, false
	}
	var c struct {
		Via       []any `json:"via"`
		Order     any   `json:"order"`
		Suggested any   `json:"suggested"`
	}
	if err := json.Unmarshal(content, &c); err != nil {
		return spaceChild{}, false
	}
	child := spaceChild{roomID: stateKey, ts: ts}
	for _, v := range c.Via {
		if s, ok := v.(string); ok && s != "" {
			child.via = append(child.via, s)
		}
	}
	if len(child.via) == 0 {
		return spaceChild{}, false
	}
	if order, ok := c.Order.(string); ok && validOrder(order) {
		child.order = order
	}
	child.suggested, _ = c.Suggested.(bool)
	return child, true
}

// parseRemoteChildren reads the children of a room another server
// described.
func parseRemoteChildren(state []SpaceChildEvent, suggestedOnly bool) []spaceChild {
	var out []spaceChild
	for _, e := range state {
		if e.Type != TypeSpaceChild {
			continue
		}
		child, ok := parseSpaceChild(e.StateKey, e.Content, e.OriginServerTS)
		if !ok || (suggestedOnly && !child.suggested) {
			continue
		}
		out = append(out, child)
	}
	sortChildren(out)
	return out
}

// validOrder reports whether an order string may sort a child: printable
// ASCII and no longer than the spec allows.
func validOrder(order string) bool {
	if order == "" || len(order) > maxOrderLength {
		return false
	}
	for i := range len(order) {
		if order[i] < 0x20 || order[i] > 0x7e {
			return false
		}
	}
	return true
}

// sortChildren puts children in the spec's order: those with an order
// first, by it; then by when the child event was sent; then by room ID,
// which only breaks ties.
func sortChildren(children []spaceChild) {
	slices.SortStableFunc(children, func(a, b spaceChild) int {
		switch {
		case a.order != "" && b.order == "":
			return -1
		case a.order == "" && b.order != "":
			return 1
		}
		return cmp.Or(
			strings.Compare(a.order, b.order),
			cmp.Compare(a.ts, b.ts),
			strings.Compare(a.roomID, b.roomID),
		)
	})
}