// Copyright (C) 2017-2023  Nexedi SA and Contributors.
//                          Kirill Smelkov <kirr@nexedi.com>
//
// This program is free software: you can Use, Study, Modify and Redistribute
// it under the terms of the GNU General Public License version 3, or (at your
// option) any later version, as published by the Free Software Foundation.
//
// You can also Link and Combine this program with other software covered by
// the terms of any of the Free Software licenses or any of the Open Source
// Initiative approved licenses and Convey the resulting work. Corresponding
// source of such a combination shall include the source code for all other
// software used.
//
// This program is distributed WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
//
// See COPYING file for full licensing terms.
// See https://www.nexedi.com/licensing for rationale and options.

package neo
// a node driven by master

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"lab.nexedi.com/kirr/go123/xnet"

	"lab.nexedi.com/kirr/neo/go/internal/log"
	"lab.nexedi.com/kirr/neo/go/internal/task"
	"lab.nexedi.com/kirr/neo/go/internal/xio"
	"lab.nexedi.com/kirr/neo/go/neo/neonet"
	"lab.nexedi.com/kirr/neo/go/neo/proto"
	"lab.nexedi.com/kirr/neo/go/neo/xneo"
)

// _MasteredNode provides base functionality of a NEO node driven by master.
//
// It connects to master, identifies to it, and handles master messages about
// δNodeTab, δPartTab and δClusterState to maintain local replica of cluster
// state. The other messages from master are passed through to _MasterLink.Recv1.
// In other words _MasteredNode installs a kind of reception pipeline in between
// master and _MasteredNode user:
//
//	                              δNodeTab
//	                              δPartTab
//	                              δClusterState
//	                              ↑
//	        _MasterLink.Recv1   ---------------
//	 user <------------------- | _MasteredNode | <- M
//	                            ---------------
//
// This pipeline is operated by TalkMaster.
// The connection to master is persisted by redial as needed.
type _MasteredNode struct {
	*xneo.Node

	// operational state is maintained by TalkMaster.
	// users retrieve it via WithOperational().
	//
	// being operational means:
	// - link to master established and is ok
	// - state of partTab/nodeTab is operational
	opMu        sync.RWMutex
	mlink       *neonet.NodeLink // link to master
	opReady	    chan struct{}    // reinitialized each time state becomes non-operational
	operational bool             // cache for Node.State.IsOperational()

	rxmFlags _MasteredNodeFlags // if e.g. δPartTab messages should be delivered to mlink.Recv1
				    // XXX just use `.myInfo.NodeType == STORAGE` instead?
}

type _MasteredNodeFlags int
const (
	// δPartTabPassThrough tells mlink.Recv1 not to filter out messages related
	// to partition table changes. When mlink.Recv1 receives such messages they
	// are already processed internally to update .state.PartTab correspondingly.
	//
	// Storage uses this mode to receive δPartTab notifications to know
	// when to persist it.
	δPartTabPassThrough _MasteredNodeFlags = 1 << iota
)

// newMasteredNode creates new _MasteredNode that connects to masters/cluster via net.
//
// Addresses of known masters are specified by masterAddrSlice.
func newMasteredNode(typ proto.NodeType, clusterName string, net xnet.Networker, masterAddrSlice []string) *_MasteredNode {
	node := &_MasteredNode{
		Node: xneo.NewNode(typ, clusterName, net, masterAddrSlice),
		opReady: make(chan struct{}),
	}

	return node
}


// TalkMaster dials master, identifies to it, and receives master notifications and requests.
//
// Notifications to node/partition tables and cluster state are automatically
// handled, while other notifications and requests are passed through to _MasterLink.Recv1.
//
// Master link is persisted by redialing as needed.
//
// f is called on every reconnection to master after identification and protocol prologue.
//
// See top-level _MasteredNode overview for details.
func (node *_MasteredNode) TalkMaster(ctx context.Context, f func(context.Context, *_MasterLink) error) (err error) {
	// start logging with initial NID (that might be temporary, and which master can tell us to change)
	ctx0 := ctx

	// When a node is created with "NewNode", we don't know yet, which of the
	// the provided master nodes is the primary master. We'll figure this out here
	// and simply start with the first node.
	maddr := node.MasterAddrSlice[0]

	defer task.Runningf(&ctx, "%s: talk master(%s)", node.MyInfo.NID, maddr)(&err)

	for {
		node.updateOperational(func() {
			node.mlink = nil
		})
		// TODO log "trying maddr as PM"
		err := node.talkMaster1(ctx, ctx0, maddr, f)
		log.Warning(ctx, err)	// XXX Warning -> Error?
		if errors.Is(err, cmdShutdown) {
			return err // M commands to shutdown
		}

		// NotPrimaryMaster -> jump to trying what it says is primary
		var notPrimaryMaster *proto.NotPrimaryMaster
		if errors.As(err, &notPrimaryMaster) {
			p := int(notPrimaryMaster.Primary)
			if !(0 <= p && p < len(notPrimaryMaster.KnownMasterList)) {
				log.Warning(ctx, "malformed NotPrimaryMaster reply - ignoring")
			} else {
				// TODO update masterRegistry from received KnownMasterList
				primary := notPrimaryMaster.KnownMasterList[p]
				maddr = primary.Address.String()
				log.Info(ctx, "switching to try %s as primary master", maddr)
			}
		}

		// TODO if err == "reject identification / protocol error" -> shutdown client?
		// TODO if err == "cannot connect" -> start trying other master nodes from MasterAddrSlice

		// exit on cancel / throttle reconnecting
		select {
		case <-ctx.Done():
			return ctx.Err()

		// XXX 1s hardcoded -> move out of here
		case <-time.After(1*time.Second):
			// ok
		}
	}
}

func (node *_MasteredNode) talkMaster1(ctx, ctxPreTalkM context.Context, maddr string, f func(context.Context, *_MasterLink) error) error {
	reqID := &proto.RequestIdentification{
		NodeType:	node.MyInfo.Type,
		NID:		node.MyInfo.NID,
		Address:	node.MyInfo.Addr,
		ClusterName:	node.ClusterName,
		IdTime:		node.MyInfo.IdTime,
		DevPath:        nil,                // XXX stub
		NewNID:         nil,	            // XXX stub
	}
	mlink, accept, err := xneo.Dial(ctx, proto.MASTER, node.Net, maddr, reqID)
	if err != nil {
		return err
	}

	return xio.WithCloseOnRetCancel(ctx, mlink, func() (err error) {
		if accept.YourNID != node.MyInfo.NID {
			log.Infof(ctx, "master %s told us to be %s", accept.MyNID, accept.YourNID)
			node.MyInfo.NID = accept.YourNID
		}

		// TODO verify Mnid = M*; our nid corresponds to our type

		// rebuild nicer task now - when we know both our and master NIDs
		// e.g. "C?: talk master(127.0.0.1:21484)" -> "C1: talk M1".
		ctx := ctxPreTalkM
		defer task.Runningf(&ctx, "%s: talk %s", accept.YourNID, accept.MyNID)(&err)

		// master pushes whole nodeTab and partTab to us right after identification

		// nodeTab
		mnt := proto.NotifyNodeInformation{}
		_, err = mlink.Expect1(&mnt)
		if err != nil {
			return fmt.Errorf("after identification: expect nodeTab: %w", err)
		}

		// partTab  (not to S and secondary M(?))
		// https://lab.nexedi.com/nexedi/neoppod/blob/v1.12-69-gd98205d0/neo/master/handlers/__init__.py#L60-67
		pt := node.State.PartTab
		if !(node.MyInfo.Type == proto.STORAGE) {
			mpt := proto.SendPartitionTable{}
			_, err = mlink.Expect1(&mpt)
			if err != nil {
				return fmt.Errorf("after identification: expect partTab: %w", err)
			}
			pt = xneo.PartTabFromMsg(&mpt)
			log.Infof(ctx, "<- parttab:\n%s", pt)
		}

		// update cluster state
		node.updateOperational(func() {
			err = node.updateNodeTab(ctx, &mnt) // the only err is cmdShutdown
			node.State.PartTab = pt
			if err == nil {
				// keep mlink=nil on shutdown so that
				// .operational does not change to y.
				node.mlink = mlink
				// TODO also set node.state.NodeTab[Mnid].link = mlink ?
			}
		})
		if err != nil {
			return err
		}

		// run user code
		return f(ctx, &_MasterLink{mlink, node})

		// TODO ? send "disconnect" on Close (? proto.CloseClient)
	})
}

// _MasterLink represents NodeLink to master with Recv1 filtered through _MasteredNode
// δstate handler.
type _MasterLink struct {
	*neonet.NodeLink
	node *_MasteredNode
}

// Recv1 receives request from master filtered through _MasteredNode δstate handler.
//
// Must be called only when master link is established - e.g. from under TalkMaster.
func (mlink *_MasterLink) Recv1(ctx context.Context) (neonet.Request, error) {
	// XXX this implementation assumes that Recv1 user is serial
	// XXX if the client code (that is run from under talkMaster1) becomes
	//     not serial and calls multiple Recv1 in parallel, the order of
	//     recvδstate calls might become different compared to the order of
	//     from-M messages -> corrupt tables.
	//
	// (for now both users - client and storage - call mlink.Recv1 serially)
	for {
		req, err := mlink.NodeLink.Recv1() // XXX cancel on ctx
		if err != nil {
			return neonet.Request{}, err
		}

		// messages for state changes are handled internally
		δstate := true
		switch req.Msg.(type) {
		default: δstate = false
		case *proto.SendPartitionTable:     // whole partTab
		case *proto.NotifyPartitionChanges: // δ(partTab)
		case *proto.NotifyNodeInformation:  // δ(nodeTab)
		case *proto.NotifyClusterState:
		}

		if δstate {
			δpt, err := mlink.node.recvδstate(ctx, req.Msg)
			passthrough := false
			if δpt && (mlink.node.rxmFlags & δPartTabPassThrough != 0) {
				passthrough = true
			}
			if !passthrough {
				req.Close()
			}
			if err != nil {
				return neonet.Request{}, err
			}
			if !passthrough {
				continue
			}
		}

		return req, nil
	}
}

//trace:event traceClusterStateChanged(cs *proto.ClusterState)

// recvδstate handles reception of δstate messages.
func (node *_MasteredNode) recvδstate(ctx context.Context, msg proto.Msg) (δpt bool, err error) {
	δpt = false

	// update .operational + notify those who was waiting for it
	node.updateOperational(func() {
		switch msg := msg.(type) {
		default:
			panic(fmt.Sprintf("unexpected message: %T", msg))

		// <- whole partTab
		case *proto.SendPartitionTable:
			δpt = true
			pt := xneo.PartTabFromMsg(msg)
			log.Infof(ctx, "<- parttab: %s", pt)
			node.State.PartTab = pt

		// <- δ(partTab)
		case *proto.NotifyPartitionChanges:
			δpt = true
			panic("TODO δ(partTab)") // XXX

		// <- δ(nodeTab)
		case *proto.NotifyNodeInformation:
			err = node.updateNodeTab(ctx, msg) // err might be command to shutdown

		case *proto.NotifyClusterState:
			log.Infof(ctx, "<- state: %s", msg.State)
			node.State.Code.Set(msg.State)
			traceClusterStateChanged(&node.State.Code)
		}
	})

	return δpt, err
}


// updateOperational calls δf under .opMu and updates .operational from current state.
// After .opMu unlock it also notifies those who was waiting for .operational to become y.
func (node *_MasteredNode) updateOperational(δf func()) {
	var opready chan struct{}

	node.opMu.Lock()
	func() {
		defer node.opMu.Unlock()
		δf()

		operational := (node.mlink != nil)  &&  node.State.IsOperational()

		//fmt.Printf("\nupdateOperatinal: %v\n", operational)
		//fmt.Printf("  mlink:   %s\n", node.mlink)
		//fmt.Printf("  state:   %s\n", node.state.Code)
		//fmt.Printf("  partTab:\n%s", node.state.PartTab)
		//fmt.Printf("  nodeTab:\n%s", node.state.NodeTab)

		if operational != node.operational {
			node.operational = operational
			if operational {
				opready = node.opReady // don't close from under opMu
			} else {
				node.opReady = make(chan struct{}) // remake for next operational waiters
			}
		}
	}()

	if opready != nil {
		//fmt.Println("updateOperational - notifying %v\n", opready)
		close(opready)
	}
}

// WithOperational runs f during when cluster state is/becomes operational.
// The cluster state is guaranteed not to change during f run.
func (node *_MasteredNode) WithOperational(ctx context.Context, f func(mlink *neonet.NodeLink, cs *xneo.ClusterState) error) error {
	for {
		node.opMu.RLock()
		if node.operational {
			//fmt.Printf("withOperation -> ready\n");
			break
		}

		ready := node.opReady
		node.opMu.RUnlock()

		//fmt.Printf("withOperational - waiting on %v\n", ready)

		select {
		case <-ctx.Done():
			return fmt.Errorf("wait operational: %w", ctx.Err())

		case <-ready:
			// ok - try to relock and read again.
		}
	}

	// node.operational=y and node.opMu is rlocked
	defer node.opMu.RUnlock()
	return f(node.mlink, &node.State)
}


var cmdShutdown = errors.New("master told us to shutdown")

// updateNodeTab applies updates to .nodeTab from message and logs changes appropriately.
// the only possible error is cmdShutdown.
// must be called under .opMu.
func (node *_MasteredNode) updateNodeTab(ctx context.Context, msg *proto.NotifyNodeInformation) error {
	// XXX msg.IdTime ?
	for _, nodeInfo := range msg.NodeList {
		log.Infof(ctx, "<- node: %v", nodeInfo)
		node.State.NodeTab.Update(nodeInfo)

		// we have to provide IdTime when requesting identification to other peers
		// (e.g. Spy checks this is what master broadcast them and if not replies "unknown by master")

		// TODO .State = DOWN -> ResetLink

		if nodeInfo.NID == node.MyInfo.NID {
			// XXX recheck locking
			node.MyInfo = nodeInfo

			// NEO/py currently employs this hack
			// FIXME -> better it be separate command and handled cleanly
			if nodeInfo.State == proto.DOWN {
				err := cmdShutdown
				log.Info(ctx, err)
				return err
			}
		}
	}

	log.Infof(ctx, "full nodetab:\n%s", node.State.NodeTab)
	return nil
}
