package cluster

import (
	"context"
	"encoding/json"
	"testing"

	"webui4frpc/internal/store"
)

// ringOf builds a 3-node ring where n1 is the leader.
func ringOf(t *testing.T) (n1, n2, n3 *Engine) {
	mk := func(id string) *Engine {
		return NewEngine(id, id+":7500", "u", "p", "0.1.2", nil,
			&fakeHandler{load: Load{MemPct: 5, NetPct: 5}},
			func(ctx context.Context, next string, tk *Token) error { return nil },
			id+":7500", true, "")
	}
	n1, n2, n3 = mk("n1"), mk("n2"), mk("n3")
	seed := State{
		LeaderID: "n1",
		Nodes: []Node{
			{ID: "n1", Addr: "n1:7500", Alive: true, IsLeader: true},
			{ID: "n2", Addr: "n2:7500", Alive: true},
			{ID: "n3", Addr: "n3:7500", Alive: true},
		},
		Topology:     map[string]*TopoEntry{},
		PendingTasks: map[string]*Task{},
	}
	for _, e := range []*Engine{n1, n2, n3} {
		// Round-trip through JSON so each engine holds its OWN Nodes slice and
		// maps. AdoptState assigns the state by value, but Nodes is a slice, so
		// handing all three the same seed would alias one backing array and make
		// a mutation on one node visible on the others — an artefact of the test
		// harness, not of the real transport (each process deserializes its own
		// copy from the token).
		blob, err := json.Marshal(seed)
		if err != nil {
			t.Fatal(err)
		}
		var own State
		if err := json.Unmarshal(blob, &own); err != nil {
			t.Fatal(err)
		}
		e.AdoptState(own)
	}
	return
}

// TestBecomeLeaderDeclinesToLiveLeader is the regression test for the live split
// brain.
//
// Leadership used to be self-declared: any node that failed to reach the leader
// it believed in promoted itself immediately. Both sides of a partition take that
// path, so two rings ran concurrently and each emptied the other's topology into
// pending — observed live as two nodes naming the same leader while marking each
// other dead, with pending=3 and topology=0 on both.
//
// A node must therefore consult the ring before promoting itself.
func TestBecomeLeaderDeclinesToLiveLeader(t *testing.T) {
	_, n2, _ := ringOf(t)
	// n3 has already promoted itself and n2 has learned that view: n1 is dead,
	// n3 is alive and leading. n2 cannot reach n3 (that is why it is asking).
	for i := range n2.state.Nodes {
		switch n2.state.Nodes[i].ID {
		case "n1":
			n2.state.Nodes[i].Alive = false
			n2.state.Nodes[i].IsLeader = false
		case "n3":
			n2.state.Nodes[i].Alive = true
			n2.state.Nodes[i].IsLeader = true
		}
	}
	n2.state.LeaderID = "n3"

	if n2.becomeLeader() {
		t.Fatal("n2 promoted itself while n3 was alive and leading in its view — " +
			"two leaders is exactly the split brain this guards against")
	}
	if n2.state.LeaderID == "n2" {
		t.Fatal("LeaderID was overwritten despite declining")
	}
	if n2.state.Nodes[0].IsLeader && n2.state.Nodes[0].ID == "n2" {
		t.Fatal("n2 marked itself leader despite declining")
	}
}

// TestBecomeLeaderPromotesWhenSoleSurvivor: the guard must not break genuine
// failover. With every other node dead, promotion has to happen.
func TestBecomeLeaderPromotesWhenSoleSurvivor(t *testing.T) {
	_, n2, _ := ringOf(t)
	for _, id := range []string{"n1", "n3"} {
		n2.state.MarkOffline(id)
	}
	if !n2.becomeLeader() {
		t.Fatal("n2 declined promotion with every other node dead — that breaks failover")
	}
	if n2.state.LeaderID != "n2" {
		t.Fatalf("leader=%q, want n2", n2.state.LeaderID)
	}
}

// TestTakeoverLeavesLiveLeadersForwardsAlone pins the second half of the fix.
//
// When a heartbeat fails, the node used to mark the leader dead and re-queue ALL
// of that leader's forwards unconditionally — even when the ring still showed a
// live leader, i.e. this node was merely the partitioned side. That is what
// emptied a working ring: both sides ended up with topology=0 and pending=3.
//
// Now the re-hang happens only if we actually take leadership.
func TestTakeoverLeavesLiveLeadersForwardsAlone(t *testing.T) {
	_, n2, _ := ringOf(t)
	// n1 (the leader n2 believes in) owns a forward.
	n2.state.AddTopology(&Task{
		ID:    "t1",
		Local: store.Local{Name: "web"}, Remote: store.Remote{Name: "frps1"},
		Link: store.Link{RemotePort: 18081},
	}, "n1")
	before := len(n2.state.TopologyList())
	pendingBefore := len(n2.state.PendingList())

	// n2 still believes n1 is the leader and cannot reach it, so it tries to take
	// over. But its state ALSO lists n3 as alive-and-leader — the view it learned
	// when n3 promoted itself. Two claims, one ring: n2 is the partitioned side
	// and must stand down.
	for i := range n2.state.Nodes {
		switch n2.state.Nodes[i].ID {
		case "n1":
			n2.state.Nodes[i].Alive = true // still marked alive; MarkOffline happens inside
			n2.state.Nodes[i].IsLeader = true
		case "n3":
			n2.state.Nodes[i].Alive = true
			n2.state.Nodes[i].IsLeader = true
		}
	}
	n2.state.LeaderID = "n1"

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	n2.takeOverFromDeadLeader(ctx)

	if n2.state.LeaderID == "n2" {
		t.Error("n2 took over while a live leader existed")
	}
	if after := len(n2.state.TopologyList()); after != before {
		t.Errorf("a declined takeover still re-queued the live leader's forwards: topology %d -> %d",
			before, after)
	}
	if p := len(n2.state.PendingList()); p != pendingBefore {
		t.Errorf("a declined takeover still pushed work into pending: %d -> %d", pendingBefore, p)
	}
}

// TestTakeoverReassignsWhenSoleSurvivor is the other half: genuine failover must
// still work, otherwise the guard would have broken the design's promise that a
// dead leader is replaced and its forwards re-hung.
func TestTakeoverReassignsWhenSoleSurvivor(t *testing.T) {
	_, n2, _ := ringOf(t)
	n2.state.AddTopology(&Task{
		ID:    "t1",
		Local: store.Local{Name: "web"}, Remote: store.Remote{Name: "frps1"},
		Link: store.Link{RemotePort: 18081},
	}, "n1") // owned by the leader n1

	// n1 is the leader and is dead; n3 is dead too. n2 is the sole survivor.
	for _, id := range []string{"n1", "n3"} {
		n2.state.MarkOffline(id)
	}
	n2.state.Nodes[0].IsLeader = false // n1 no longer claims leadership

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	n2.takeOverFromDeadLeader(ctx)

	if n2.state.LeaderID != "n2" {
		t.Fatalf("n2 did not take over as the sole survivor, leader=%q", n2.state.LeaderID)
	}
	if len(n2.state.TopologyList()) != 0 {
		t.Error("the dead leader's forwards should have been re-hung")
	}
	found := false
	for _, tk := range n2.state.PendingList() {
		if tk.Local.Name == "web" {
			found = true
		}
	}
	if !found {
		t.Error("the dead leader's forward was not re-queued for another node")
	}
}