package main

import (
	"fmt"
	"math/rand"
	"sync"
	"time"
)

type ServerState string

const (
	Leader    ServerState = "Leader"
	Follower  ServerState = "Follower"
	Candidate ServerState = "Candidate"
)

type LogEntry struct {
	Command string
	Term    int
}

type Node struct {
	ID          int
	State       ServerState
	CurrentTerm int
	VotedFor    *int
	Log         []LogEntry
	IsAlive     bool
	Votes       int
}

type Cluster struct {
	Nodes      []*Node
	LeaderID   *int
	mu         sync.Mutex
	ElectionCh chan int
}

func NewCluster(n int) *Cluster {
	nodes := make([]*Node, n)
	for i := range nodes {
		nodes[i] = &Node{
			ID:          i + 1,
			State:       Follower,
			CurrentTerm: 0,
			VotedFor:    nil,
			Log:         []LogEntry{},
			IsAlive:     true,
			Votes:       0,
		}
	}
	return &Cluster{Nodes: nodes, LeaderID: nil, ElectionCh: make(chan int)}
}

func (c *Cluster) StartElection(node *Node) {
	c.mu.Lock()
	defer c.mu.Unlock()

	node.State = Candidate
	node.CurrentTerm++
	node.VotedFor = &node.ID
	node.Votes = 1 // Vote for self
	fmt.Printf("Node %d started election for term %d\n", node.ID, node.CurrentTerm)

	for _, peer := range c.Nodes {
		if peer.ID != node.ID && peer.IsAlive && peer.CurrentTerm <= node.CurrentTerm {
			peer.VotedFor = &node.ID
			node.Votes++
		}
	}

	if node.Votes > len(c.Nodes)/2 {
		c.LeaderID = &node.ID
		node.State = Leader
		fmt.Printf("Node %d became the leader for term %d\n", node.ID, node.CurrentTerm)
		c.ElectionCh <- node.ID
	}
}

func (c *Cluster) Heartbeat(node *Node) {
	for {
		select {
		case leaderID := <-c.ElectionCh:
			if node.ID != leaderID && node.IsAlive {
				fmt.Printf("Node %d acknowledges leader %d\n", node.ID, leaderID)
				node.State = Follower
				node.VotedFor = nil
			}
		case <-time.After(2 * time.Second):
			if node.State == Leader && node.IsAlive {
				fmt.Printf("Leader %d sending heartbeat\n", node.ID)
			} else if node.State != Leader && node.IsAlive {
				fmt.Printf("Node %d has not received heartbeat, starting election\n", node.ID)
				c.StartElection(node)
			}
		}
	}
}

func (c *Cluster) Simulate(duration time.Duration) {
	fmt.Println("Starting simulation...")
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()
	end := time.After(duration)

	for {
		select {
		case <-ticker.C:
			c.mu.Lock()
			node := c.Nodes[rand.Intn(len(c.Nodes))]
			node.IsAlive = !node.IsAlive // Toggle alive status for simulation
			status := "failed"
			if node.IsAlive {
				status = "recovered"
			}
			fmt.Printf("Node %d has %s\n", node.ID, status)
			c.mu.Unlock()
		case <-end:
			fmt.Println("Simulation ended.")
			return
		}
	}
}

func main() {
	rand.Seed(time.Now().UnixNano())
	cluster := NewCluster(5)

	// Launch Heartbeat and Election goroutines
	for _, node := range cluster.Nodes {
		go cluster.Heartbeat(node)
	}

	// Run the simulation for 10 seconds
	cluster.Simulate(10 * time.Second)
}