// NewPropagationFunc registers a new protocol name with the context c and will
// set f as handler for every new instance of that protocol.
func NewPropagationFunc(c propagationContext, name string, f PropagationStore) (PropagationFunc, error) {
	pid, err := c.ProtocolRegister(name, func(n *onet.TreeNodeInstance) (onet.ProtocolInstance, error) {
		p := &Propagate{
			sd:               &PropagateSendData{[]byte{}, 1000},
			TreeNodeInstance: n,
			received:         0,
			subtreeCount:     n.TreeNode().SubtreeCount(),
			onData:           f,
		}
		for _, h := range []interface{}{&p.ChannelSD, &p.ChannelReply} {
			if err := p.RegisterChannel(h); err != nil {
				return nil, err
			}
		}
		return p, nil
	})
	log.Lvl3("Registering new propagation for", c.ServerIdentity(),
		name, pid)
	return func(el *onet.Roster, msg network.Body, msec int) (int, error) {
		tree := el.GenerateNaryTreeWithRoot(8, c.ServerIdentity())
		log.Lvl3(el.List[0].Address, "Starting to propagate", reflect.TypeOf(msg))
		pi, err := c.CreateProtocolOnet(name, tree)
		if err != nil {
			return -1, err
		}
		return propagateStartAndWait(pi, msg, msec, f)
	}, err
}

// Dispatch listens for all channels and waits for a timeout in case nothing
// happens for a certain duration
func (p *ProtocolCount) Dispatch() error {
	running := true
	for running {
		log.Lvl3(p.Info(), "waiting for message during", p.Timeout())
		select {
		case pc := <-p.PrepareCountChan:
			log.Lvl3(p.Info(), "received from", pc.TreeNode.ServerIdentity.Address,
				pc.Timeout)
			p.SetTimeout(pc.Timeout)
			p.FuncPC()
		case c := <-p.CountChan:
			p.FuncC(c)
			running = false
		case _ = <-p.NodeIsUpChan:
			if p.Parent() != nil {
				err := p.SendTo(p.Parent(), &NodeIsUp{})
				if err != nil {
					log.Error(p.Info(), "couldn't send to parent",
						p.Parent().Name(), err)
				}
			} else {
				p.Replies++
			}
		case <-time.After(time.Duration(p.Timeout()) * time.Millisecond):
			log.Lvl3(p.Info(), "timed out while waiting for", p.Timeout())
			if p.IsRoot() {
				log.Lvl2("Didn't get all children in time:", p.Replies)
				p.Count <- p.Replies
				running = false
			}
		}
	}
	p.Done()
	return nil
}

// handleCommit receives commit messages and signal the end if it received
// enough of it.
func (p *Protocol) handleCommit(com *Commit) {
	if p.state != stateCommit {
		//	log.Lvl3(p.Name(), "STORE handle commit packet")
		p.tempCommitMsg = append(p.tempCommitMsg, com)
		return
	}
	// finish after threshold of Commit msgs
	p.commitMsgCount++
	log.Lvl4(p.Name(), "----------------\nWe got", p.commitMsgCount,
		"COMMIT msgs and threshold is", p.threshold)
	if p.IsRoot() {
		log.Lvl4("Leader got ", p.commitMsgCount)
	}
	if p.commitMsgCount >= p.threshold {
		p.state = stateFinished
		// reset counter
		p.commitMsgCount = 0
		log.Lvl3(p.Name(), "Threshold reached: We are done... CONSENSUS")
		if p.IsRoot() && p.onDoneCB != nil {
			log.Lvl3(p.Name(), "We are root and threshold reached: return to the simulation.")
			p.onDoneCB()
			p.finish()
		}
		return
	}
}

func runProtocolOnceGo(nbrHosts int, name string, refuseCount int,
	succeed bool) error {
	log.Lvl2("Running BFTCoSi with", nbrHosts, "hosts")
	local := onet.NewLocalTest()
	defer local.CloseAll()
	_, _, tree := local.GenBigTree(nbrHosts, nbrHosts, 2, true)
	log.Lvl3("Tree is:", tree.Dump())

	done := make(chan bool)
	// create the message we want to sign for this round
	msg := []byte("Hello BFTCoSi")

	// Start the protocol
	node, err := local.CreateProtocol(name, tree)
	if err != nil {
		return errors.New("Couldn't create new node: " + err.Error())
	}

	// Register the function generating the protocol instance
	var root *ProtocolBFTCoSi
	root = node.(*ProtocolBFTCoSi)
	root.Msg = msg
	cMux.Lock()
	counter := &Counter{refuseCount: refuseCount}
	counters.add(counter)
	root.Data = []byte(strconv.Itoa(counters.size() - 1))
	log.Lvl3("Added counter", counters.size()-1, refuseCount)
	cMux.Unlock()
	log.ErrFatal(err)
	// function that will be called when protocol is finished by the root
	root.RegisterOnDone(func() {
		done <- true
	})
	go node.Start()
	log.Lvl1("Launched protocol")
	// are we done yet?
	wait := time.Second * 60
	select {
	case <-done:
		counter.Lock()
		if counter.veriCount != nbrHosts {
			return errors.New("Each host should have called verification.")
		}
		// if assert refuses we don't care for unlocking (t.Refuse)
		counter.Unlock()
		sig := root.Signature()
		err := sig.Verify(root.Suite(), root.Roster().Publics())
		if succeed && err != nil {
			return fmt.Errorf("%s Verification of the signature refused: %s - %+v", root.Name(), err.Error(), sig.Sig)
		}
		if !succeed && err == nil {
			return fmt.Errorf("%s: Shouldn't have succeeded for %d hosts, but signed for count: %d",
				root.Name(), nbrHosts, refuseCount)
		}
	case <-time.After(wait):
		log.Lvl1("Going to break because of timeout")
		return errors.New("Waited " + wait.String() + " for BFTCoSi to finish ...")
	}
	return nil
}

// handleAnnounce receive the announcement from another node
// it reply with an ACK.
func (b *Broadcast) handleContactNodes(msg struct {
	*onet.TreeNode
	ContactNodes
}) error {
	log.Lvl3(b.Info(), "Received message from", msg.TreeNode.String())
	if msg.TreeNode.ID == b.Root().ID {
		b.repliesLeft = len(b.Tree().List()) - b.tnIndex - 1
		if b.repliesLeft == 0 {
			log.Lvl3("Won't contact anybody - finishing")
			b.SendTo(b.Root(), &Done{})
			return nil
		}
		log.Lvl3(b.Info(), "Contacting nodes:", b.repliesLeft)
		// Connect to all nodes that are later in the TreeNodeList, but only if
		// the message comes from root
		for _, tn := range b.Tree().List()[b.tnIndex+1:] {
			log.Lvl3("Connecting to", tn.String())
			err := b.SendTo(tn, &ContactNodes{})
			if err != nil {
				return nil
			}
		}
	} else {
		// Tell the caller we're done
		log.Lvl3("Sending back to", msg.TreeNode.ServerIdentity.String())
		b.SendTo(msg.TreeNode, &Done{})
	}
	return nil
}

// waitResponseVerification waits till the end of the verification and returns
// the BFTCoSiResponse along with the flag:
// true => no exception, the verification is correct
// false => exception, the verification failed
func (bft *ProtocolBFTCoSi) waitResponseVerification() (*Response, bool) {
	log.Lvl3(bft.Name(), "Waiting for response verification:")
	// wait the verification
	verified := <-bft.verifyChan

	resp, err := bft.prepare.Response(bft.tempPrepareResponse)
	if err != nil {
		return nil, false
	}

	if !verified {
		// Add our exception
		bft.tempExceptions = append(bft.tempExceptions, Exception{
			Index:      bft.index,
			Commitment: bft.prepare.GetCommitment(),
		})
		// Don't include our response!
		resp = bft.Suite().Scalar().Set(resp).Sub(resp, bft.prepare.GetResponse())
		log.Lvl3(bft.Name(), "Response verification: failed")
	}

	r := &Response{
		TYPE:       RoundPrepare,
		Exceptions: bft.tempExceptions,
		Response:   resp,
	}

	log.Lvl3(bft.Name(), "Response verification:", verified)
	return r, verified
}

func TestThreshold(t *testing.T) {
	const TestProtocolName = "DummyBFTCoSiThr"

	// Register test protocol using BFTCoSi
	onet.GlobalProtocolRegister(TestProtocolName, func(n *onet.TreeNodeInstance) (onet.ProtocolInstance, error) {
		return NewBFTCoSiProtocol(n, verify)
	})

	local := onet.NewLocalTest()
	defer local.CloseAll()
	tests := []struct{ h, t int }{
		{1, 1},
		{2, 2},
		{3, 2},
		{4, 3},
		{5, 4},
		{6, 4},
	}
	for _, s := range tests {
		hosts, thr := s.h, s.t
		log.Lvl3("Hosts is", hosts)
		_, _, tree := local.GenBigTree(hosts, hosts, 2, true)
		log.Lvl3("Tree is:", tree.Dump())

		// Start the protocol
		node, err := local.CreateProtocol(TestProtocolName, tree)
		log.ErrFatal(err)
		bc := node.(*ProtocolBFTCoSi)
		assert.Equal(t, thr, bc.threshold, "hosts was %d", hosts)
		local.CloseAll()
	}
}

// CreateIdentity will register a new SkipChain and add it to our list of
// managed identities.
func (s *Service) CreateIdentity(ai *CreateIdentity) (network.Body, onet.ClientError) {
	log.Lvlf3("%s Creating new identity with config %+v", s, ai.Config)
	ids := &Storage{
		Latest: ai.Config,
	}
	log.Lvl3("Creating Root-skipchain")
	var cerr onet.ClientError
	ids.Root, cerr = s.skipchain.CreateRoster(ai.Roster, 2, 10,
		skipchain.VerifyNone, nil)
	if cerr != nil {
		return nil, cerr
	}
	log.Lvl3("Creating Data-skipchain")
	ids.Root, ids.Data, cerr = s.skipchain.CreateData(ids.Root, 2, 10,
		skipchain.VerifyNone, ai.Config)
	if cerr != nil {
		return nil, cerr
	}

	roster := ids.Root.Roster
	replies, err := s.propagateIdentity(roster, &PropagateIdentity{ids}, propagateTimeout)
	if err != nil {
		return nil, onet.NewClientErrorCode(ErrorOnet, err.Error())
	}
	if replies != len(roster.List) {
		log.Warn("Did only get", replies, "out of", len(roster.List))
	}
	log.Lvlf2("New chain is\n%x", []byte(ids.Data.Hash))
	s.save()

	return &CreateIdentityReply{
		Root: ids.Root,
		Data: ids.Data,
	}, nil
}

// nodeDone is either called by the end of EndProtocol or by the end of the
// response phase of the commit round.
func (bz *ByzCoin) nodeDone() bool {
	log.Lvl3(bz.Name(), "nodeDone()      ----- ")
	bz.doneProcessing <- true
	log.Lvl3(bz.Name(), "nodeDone()      +++++  ", bz.onDoneCallback)
	if bz.onDoneCallback != nil {
		bz.onDoneCallback()
	}
	return true
}

// FuncPrepareClose sends a `PrepareClose`-message down the tree.
func (p *ProtocolCloseAll) FuncPrepareClose(pc PrepareCloseMsg) error {
	log.Lvl3(pc.ServerIdentity.Address, "sent PrepClose to", p.ServerIdentity().Address)
	if !p.IsLeaf() {
		for _, c := range p.Children() {
			err := p.SendTo(c, &PrepareClose{})
			log.Lvl3(p.ServerIdentity().Address, "sends to", c.ServerIdentity.Address, "(err=", err, ")")
		}
	} else {
		p.FuncClose(nil)
	}
	return nil
}

// Dispatch can handle timeouts
func (p *Propagate) Dispatch() error {
	process := true
	log.Lvl4(p.ServerIdentity())
	for process {
		p.Lock()
		timeout := time.Millisecond * time.Duration(p.sd.Msec)
		p.Unlock()
		select {
		case msg := <-p.ChannelSD:
			log.Lvl3(p.ServerIdentity(), "Got data from", msg.ServerIdentity, "and setting timeout to", msg.Msec)
			p.sd.Msec = msg.Msec
			if p.onData != nil {
				_, netMsg, err := network.UnmarshalRegistered(msg.Data)
				if err == nil {
					p.onData(netMsg)
				}
			}
			if !p.IsRoot() {
				log.Lvl3(p.ServerIdentity(), "Sending to parent")
				p.SendToParent(&PropagateReply{})
			}
			if p.IsLeaf() {
				process = false
			} else {
				log.Lvl3(p.ServerIdentity(), "Sending to children")
				p.SendToChildren(&msg.PropagateSendData)
			}
		case <-p.ChannelReply:
			p.received++
			log.Lvl4(p.ServerIdentity(), "received:", p.received, p.subtreeCount)
			if !p.IsRoot() {
				p.SendToParent(&PropagateReply{})
			}
			if p.received == p.subtreeCount {
				process = false
			}
		case <-time.After(timeout):
			_, a, err := network.UnmarshalRegistered(p.sd.Data)
			log.Fatalf("Timeout of %s reached. %v %s", timeout, a, err)
			process = false
		}
	}
	if p.IsRoot() {
		if p.onDoneCb != nil {
			p.onDoneCb(p.received + 1)
		}
	}
	p.Done()
	return nil
}

// network will contact all cothorities in the group-file and print
// the status-report of each one.
func network(c *cli.Context) error {
	groupToml := c.GlobalString("g")
	el, err := readGroup(groupToml)
	log.ErrFatal(err, "Couldn't Read File")
	log.Lvl3(el)
	cl := status.NewClient()
	for i := 0; i < len(el.List); i++ {
		log.Lvl3(el.List[i])
		sr, _ := cl.Request(el.List[i])
		printConn(sr)
		log.Lvl3(cl)
	}
	return nil
}

// HandleReply is the message going up the tree and holding a counter
// to verify the number of nodes.
func (p *ProtocolExampleHandlers) HandleReply(reply []StructReply) error {
	children := 1
	for _, c := range reply {
		children += c.ChildrenCount
	}
	log.Lvl3(p.ServerIdentity().Address, "is done with total of", children)
	if !p.IsRoot() {
		log.Lvl3("Sending to parent")
		return p.SendTo(p.Parent(), &Reply{children})
	}
	log.Lvl3("Root-node is done - nbr of children found:", children)
	p.ChildCount <- children
	return nil
}

// listen will select on the differents channels
func (nt *Ntree) listen() {
	for {
		select {
		// Dispatch the block through the whole tree
		case msg := <-nt.announceChan:
			log.Lvl3(nt.Name(), "Received Block announcement")
			nt.block = msg.BlockAnnounce.Block
			// verify the block
			go byzcoin.VerifyBlock(nt.block, "", "", nt.verifyBlockChan)
			if nt.IsLeaf() {
				nt.startBlockSignature()
				continue
			}
			for _, tn := range nt.Children() {
				err := nt.SendTo(tn, &msg.BlockAnnounce)
				if err != nil {
					log.Error(nt.Name(),
						"couldn't send to", tn.Name(),
						err)
				}
			}
			// generate your own signature / exception and pass that up to the
			// root
		case msg := <-nt.blockSignatureChan:
			nt.handleBlockSignature(&msg.NaiveBlockSignature)
			// Dispatch the signature + expcetion made before through the whole
			// tree
		case msg := <-nt.roundSignatureRequestChan:
			log.Lvl3(nt.Name(), " Signature Request Received")
			go nt.verifySignatureRequest(&msg.RoundSignatureRequest)

			if nt.IsLeaf() {
				nt.startSignatureResponse()
				continue
			}

			for _, tn := range nt.Children() {
				err := nt.SendTo(tn, &msg.RoundSignatureRequest)
				if err != nil {
					log.Error(nt.Name(), "couldn't sent to",
						tn.Name(), err)
				}
			}
			// Decide if we want to sign this or not
		case msg := <-nt.roundSignatureResponseChan:
			nt.handleRoundSignatureResponse(&msg.RoundSignatureResponse)
		}
	}
}

// Start the last phase : send up the final signature
func (nt *Ntree) startSignatureResponse() {
	log.Lvl3(nt.Name(), "Start Signature Response phase")
	nt.computeSignatureResponse()
	if err := nt.SendTo(nt.Parent(), nt.tempSignatureResponse); err != nil {
		log.Error(err)
	}
}

// SignatureRequest treats external request to this service.
func (cs *CoSi) SignatureRequest(req *SignatureRequest) (network.Body, onet.ClientError) {
	if req.Roster.ID == onet.RosterID(uuid.Nil) {
		req.Roster.ID = onet.RosterID(uuid.NewV4())
	}
	tree := req.Roster.GenerateBinaryTree()
	tni := cs.NewTreeNodeInstance(tree, tree.Root, cosi.Name)
	pi, err := cosi.NewProtocol(tni)
	if err != nil {
		return nil, onet.NewClientErrorCode(4100, "Couldn't make new protocol: "+err.Error())
	}
	cs.RegisterProtocolInstance(pi)
	pcosi := pi.(*cosi.CoSi)
	pcosi.SigningMessage(req.Message)
	h, err := crypto.HashBytes(network.Suite.Hash(), req.Message)
	if err != nil {
		return nil, onet.NewClientErrorCode(4101, "Couldn't hash message: "+err.Error())
	}
	response := make(chan []byte)
	pcosi.RegisterSignatureHook(func(sig []byte) {
		response <- sig
	})
	log.Lvl3("Cosi Service starting up root protocol")
	go pi.Dispatch()
	go pi.Start()
	sig := <-response
	if log.DebugVisible() > 1 {
		fmt.Printf("%s: Signed a message.\n", time.Now().Format("Mon Jan 2 15:04:05 -0700 MST 2006"))
	}
	return &SignatureResponse{
		Hash:      h,
		Signature: sig,
	}, nil
}

// startPrepareChallenge create the challenge and send its down the tree
func (bz *ByzCoin) startChallengePrepare() error {
	// make the challenge out of it
	trblock := bz.tempBlock
	marshalled, err := json.Marshal(trblock)
	if err != nil {
		return err
	}
	ch, err := bz.prepare.CreateChallenge(marshalled)
	if err != nil {
		return err
	}
	bizChal := &ChallengePrepare{
		TYPE:      RoundPrepare,
		Challenge: ch,
		TrBlock:   trblock,
	}

	go VerifyBlock(bz.tempBlock, bz.lastBlock, bz.lastKeyBlock, bz.verifyBlockChan)
	log.Lvl3(bz.Name(), "ByzCoin Start Challenge PREPARE")
	// send to children
	for _, tn := range bz.Children() {
		err = bz.SendTo(tn, bizChal)
	}
	return err
}

func (bz *ByzCoin) handleResponsePrepare(bzr *Response) error {
	// check if we have enough
	bz.tprMut.Lock()
	bz.tempPrepareResponse = append(bz.tempPrepareResponse, bzr.Response)
	if len(bz.tempPrepareResponse) < len(bz.Children()) {
		bz.tprMut.Unlock()
		return nil
	}

	// wait for verification
	bzrReturn, ok := bz.waitResponseVerification()
	if ok {
		// append response
		resp, err := bz.prepare.Response(bz.tempPrepareResponse)
		bz.tprMut.Unlock()
		if err != nil {
			return err
		}
		bzrReturn.Response = resp
	} else {
		bz.tprMut.Unlock()
	}

	log.Lvl3("ByzCoin Handle Response PREPARE")
	// if I'm root, we are finished, let's notify the "commit" round
	if bz.IsRoot() {
		// notify listeners (simulation) we finished
		if bz.onResponsePrepareDone != nil {
			bz.onResponsePrepareDone()
		}
		return bz.startChallengeCommit()
	}
	// send up
	return bz.SendTo(bz.Parent(), bzrReturn)
}

// Start will contact everyone and make the connections
func (b *Broadcast) Start() error {
	n := len(b.Tree().List())
	b.repliesLeft = n * (n - 1) / 2
	log.Lvl3(b.Name(), "Sending announce to", b.repliesLeft, "nodes")
	b.SendTo(b.Root(), &ContactNodes{})
	return nil
}

// Servers contacts all servers in the entity-list and then makes checks
// on each pair. If server-descriptions are available, it will print them
// along with the IP-address of the server.
// In case a server doesn't reply in time or there is an error in the
// signature, an error is returned.
func Servers(g *config.Group) error {
	success := true
	// First check all servers individually
	for _, e := range g.Roster.List {
		desc := []string{"none", "none"}
		if d := g.GetDescription(e); d != "" {
			desc = []string{d, d}
		}
		el := onet.NewRoster([]*network.ServerIdentity{e})
		success = checkList(el, desc) == nil && success
	}
	if len(g.Roster.List) > 1 {
		// Then check pairs of servers
		for i, first := range g.Roster.List {
			for _, second := range g.Roster.List[i+1:] {
				log.Lvl3("Testing connection between", first, second)
				desc := []string{"none", "none"}
				if d1 := g.GetDescription(first); d1 != "" {
					desc = []string{d1, g.GetDescription(second)}
				}
				es := []*network.ServerIdentity{first, second}
				success = checkList(onet.NewRoster(es), desc) == nil && success
				es[0], es[1] = es[1], es[0]
				desc[0], desc[1] = desc[1], desc[0]
				success = checkList(onet.NewRoster(es), desc) == nil && success
			}
		}
	}

	if !success {
		return errors.New("At least one of the tests failed")
	}
	return nil
}