// Copyright 2021 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package storage

import (
	"encoding/json"
	"sort"

	"go.uber.org/zap"

	"go.etcd.io/etcd/client/pkg/v3/types"
	"go.etcd.io/etcd/pkg/v3/pbutil"
	"go.etcd.io/etcd/server/v3/etcdserver/api/membership"
	"go.etcd.io/raft/v3/raftpb"
)

// CreateConfigChangeEnts creates a series of Raft entries (i.e.
// EntryConfChange) to remove the set of given IDs from the cluster. The ID
// `self` is _not_ removed, even if present in the set.
// If `self` is not inside the given ids, it creates a Raft entry to add a
// default member with the given `self`.
func CreateConfigChangeEnts(lg *zap.Logger, ids []uint64, self uint64, term, index uint64) []*raftpb.Entry {
	found := false
	for _, id := range ids {
		if id == self {
			found = true
		}
	}

	var ents []*raftpb.Entry
	next := index + 1

	// NB: always add self first, then remove other nodes. Raft will panic if the
	// set of voters ever becomes empty.
	if !found {
		m := membership.Member{
			ID:             types.ID(self),
			RaftAttributes: membership.RaftAttributes{PeerURLs: []string{"http://localhost:2380"}},
		}
		ctx, err := json.Marshal(m)
		if err != nil {
			lg.Panic("failed to marshal member", zap.Error(err))
		}
		cc := &raftpb.ConfChange{
			Type:    raftpb.ConfChangeAddNode.Enum(),
			NodeId:  &self,
			Context: ctx,
		}
		idx := next
		e := &raftpb.Entry{
			Type:  raftpb.EntryConfChange.Enum(),
			Data:  pbutil.MustMarshalMessage(cc),
			Term:  new(term),
			Index: new(idx),
		}
		ents = append(ents, e)
		next++
	}

	for _, id := range ids {
		if id == self {
			continue
		}
		cc := &raftpb.ConfChange{
			Type:   raftpb.ConfChangeRemoveNode.Enum(),
			NodeId: new(uint64),
		}
		*cc.NodeId = id
		idx := next
		e := &raftpb.Entry{
			Type:  raftpb.EntryConfChange.Enum(),
			Data:  pbutil.MustMarshalMessage(cc),
			Term:  new(term),
			Index: new(idx),
		}
		ents = append(ents, e)
		next++
	}

	return ents
}

// GetEffectiveNodeIDsFromWALEntries returns an ordered set of IDs included in the given snapshot and
// the entries. The given snapshot/entries can contain three kinds of
// ID-related entry:
// - ConfChangeAddNode, in which case the contained ID will Be added into the set.
// - ConfChangeRemoveNode, in which case the contained ID will Be removed from the set.
// - ConfChangeAddLearnerNode, in which the contained ID will Be added into the set.
func GetEffectiveNodeIDsFromWALEntries(lg *zap.Logger, snap *raftpb.Snapshot, ents []*raftpb.Entry) []uint64 {
	ids := make(map[uint64]bool)
	if snap.GetMetadata().GetConfState() != nil {
		for _, id := range snap.GetMetadata().GetConfState().GetVoters() {
			ids[id] = true
		}
		for _, id := range snap.GetMetadata().GetConfState().GetLearners() {
			ids[id] = true
		}
	}
	for i := range ents {
		e := ents[i]
		if e.GetType() != raftpb.EntryConfChange {
			continue
		}
		var cc raftpb.ConfChange
		pbutil.MustUnmarshalMessage(&cc, e.Data)
		switch cc.GetType() {
		case raftpb.ConfChangeAddLearnerNode:
			ids[cc.GetNodeId()] = true
		case raftpb.ConfChangeAddNode:
			ids[cc.GetNodeId()] = true
		case raftpb.ConfChangeRemoveNode:
			delete(ids, cc.GetNodeId())
		case raftpb.ConfChangeUpdateNode:
			// do nothing
		default:
			lg.Panic("unknown ConfChange Type", zap.String("type", cc.GetType().String()))
		}
	}
	sids := make(types.Uint64Slice, 0, len(ids))
	for id := range ids {
		sids = append(sids, id)
	}
	sort.Sort(sids)
	return sids
}
