package main

import (
	"bytes"
	"context"
	"encoding/gob"
	"log/slog"
	"math/rand"
	"net"
	"sync"
	"time"
)

// this file handles gossip protocol from MP2
// thread 1 - handles receiving gossip
// thread 2 - handles sending gossips and updating heart beat
// thread 3 - handles failure detection

type Member_ID struct {
	IP   string
	Port int
	Time time.Time
}

type MemberInfo struct {
	HeartBeat          int
	State              string
	local_failed_since *time.Time
	local_last_seen    *time.Time
}

type MembershipService struct {
	MembersList  map[Member_ID]MemberInfo
	SelfID       Member_ID
	deadSessions map[Member_ID]struct{}
	mutex        sync.RWMutex
}

type GossipMessage struct {
	Type       string
	Senders_id Member_ID
	Members    map[Member_ID]MemberInfo
}

// One helper for gob encoding and UDP sending
func send_gossip_message(
	socket *net.UDPConn,
	addr *net.UDPAddr,
	gossip_msg GossipMessage,
) error {
	var buffer bytes.Buffer

	if err := gob.NewEncoder(&buffer).Encode(gossip_msg); err != nil {
		return err
	}

	_, err := socket.WriteToUDP(buffer.Bytes(), addr)
	return err
}

func (m *MembershipService) get_current_members() map[Member_ID]MemberInfo {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	// Create a copy of the members_list to avoid race conditions
	current_members := make(map[Member_ID]MemberInfo)

	for id, info := range m.MembersList {
		current_members[id] = info
	}

	return current_members
}

func (m *MembershipService) get_self_id() Member_ID {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	return m.SelfID
}

func get_all_member_ids(
	members map[Member_ID]MemberInfo,
	self Member_ID,
) []Member_ID {
	member_ids := make([]Member_ID, 0, len(members))

	for id, _ := range members {
		// Exclude self and DOWN nodes
		if id != self {
			member_ids = append(member_ids, id)
		}
	}

	return member_ids
}

func (m *MembershipService) member_init(ip string, port int) {
	now := time.Now()

	// Round(0) removes monotonic clock data for gob map keys
	my_id := Member_ID{ip, port, now.UTC().Round(0)}
	members_info := MemberInfo{0, "alive", nil, &now}

	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.SelfID = my_id
	m.MembersList[my_id] = members_info
}

// Add a new process session.
// Caller must hold m.mutex.Lock().
func (m *MembershipService) add_new_member(
	id Member_ID,
	info MemberInfo,
	now time.Time,
) bool {
	// Never accept a session already declared dead
	if _, dead := m.deadSessions[id]; dead {
		return false
	}

	// Check whether we already know a newer session
	for old_id := range m.MembersList {
		if old_id.IP == id.IP && old_id.Port == id.Port {
			if !id.Time.After(old_id.Time) {
				return false
			}
		}
	}

	// New process lifetime: retire older sessions
	for old_id := range m.MembersList {
		if old_id.IP == id.IP && old_id.Port == id.Port {
			m.deadSessions[old_id] = struct{}{}
			delete(m.MembersList, old_id)
		}
	}

	info.local_last_seen = &now
	info.local_failed_since = nil

	if info.State == "down" {
		info.local_failed_since = &now
		m.deadSessions[id] = struct{}{}
	}

	m.MembersList[id] = info
	return true
}

func (m *MembershipService) handle_join(
	new_member Member_ID,
	logger *slog.Logger,
	socket *net.UDPConn,
) {
	if net.ParseIP(new_member.IP) == nil {
		logger.Warn("Invalid JOIN member", "member", new_member)
		return
	}

	m.mutex.Lock()

	// Add only if this is a new process session.
	// Repeated JOIN does not reset heartbeat tracking.
	now := time.Now()
	m.add_new_member(
		new_member,
		MemberInfo{0, "alive", nil, nil},
		now,
	)

	m.mutex.Unlock()

	// ACK with current membership list
	gossip_msg := GossipMessage{
		Type:       "join_ack",
		Senders_id: m.get_self_id(),
		Members:    m.get_current_members(),
	}

	err := send_gossip_message(
		socket,
		&net.UDPAddr{
			IP:   net.ParseIP(new_member.IP),
			Port: new_member.Port,
		},
		gossip_msg,
	)

	if err != nil {
		logger.Warn("Error sending JOIN acknowledgment", "error", err)
		return
	}

	logger.Info("JOIN acknowledgment sent", "member", new_member)
}

func (m *MembershipService) handle_gossip_receive(
	ctx context.Context,
	socket *net.UDPConn,
	logger *slog.Logger,
	config Config,
	wg *sync.WaitGroup,
) {
	defer wg.Done()

	buffer := make([]byte, 65535)

	// Introducer doesn't need to JOIN
	is_introducer := config.introducer_ip == config.my_ip.String()
	joined := is_introducer
	last_join_time := time.Time{}

	for {
		if ctx.Err() != nil {
			return
		}

		if !joined && time.Since(last_join_time) >= config.join_retry_interval {
			// Send JOIN to introducer
			join_msg := GossipMessage{
				Type:       "join",
				Senders_id: m.get_self_id(),
				Members:    map[Member_ID]MemberInfo{},
			}

			introducer_addr := &net.UDPAddr{
				IP:   net.ParseIP(config.introducer_ip),
				Port: config.port,
			}

			if err := send_gossip_message(socket, introducer_addr, join_msg); err != nil {
				logger.Warn("Error sending JOIN", "error", err)
			}

			last_join_time = time.Now()
		}

		socket.SetReadDeadline(time.Now().Add(500 * time.Millisecond))

		n, _, err := socket.ReadFromUDP(buffer)

		if err != nil {
			if ctx.Err() != nil {
				return
			}

			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue
			}

			logger.Warn("Error reading UDP socket", "error", err)
			continue
		}

		var gossip_msg GossipMessage

		// Decode gob message
		err = gob.NewDecoder(
			bytes.NewReader(buffer[:n]),
		).Decode(&gossip_msg)

		if err != nil {
			logger.Warn("Invalid gossip packet", "error", err)
			continue
		}

		switch gossip_msg.Type {

		case "gossip":
			// Merge incoming membership records
			// and rejoin if our own session was falsely marked DOWN
			if m.handle_gossip(gossip_msg.Members, logger) && !is_introducer {
				joined = false
				last_join_time = time.Time{}
			}

		case "join":
			if is_introducer {
				m.handle_join(gossip_msg.Senders_id, logger, socket)
			}

		case "join_ack":
			// Ignore delayed ACKs for an older local session
			self_id := m.get_self_id()
			if _, exists := gossip_msg.Members[self_id]; !exists {
				continue
			}

			restarted := m.handle_gossip(gossip_msg.Members, logger)

			if restarted {
				joined = false
				last_join_time = time.Time{}
				continue
			}

			if !joined {
				logger.Info("Successfully joined membership group")
			}

			joined = true

		default:
			logger.Warn("Unknown gossip message", "type", gossip_msg.Type)
		}
	}
}

func select_peers(
	current_members map[Member_ID]MemberInfo,
	fanout int,
	self Member_ID,
) []Member_ID {
	member_ids := get_all_member_ids(current_members, self)

	if fanout <= 0 {
		return nil
	}

	if len(member_ids) <= fanout {
		return member_ids
	}

	selected := make([]Member_ID, 0, fanout)
	perm := rand.Perm(len(member_ids))

	for i := 0; i < fanout; i++ {
		selected = append(selected, member_ids[perm[i]])
	}

	return selected
}

func (m *MembershipService) handle_gossip_send(
	ctx context.Context,
	socket *net.UDPConn,
	logger *slog.Logger,
	config Config,
	wg *sync.WaitGroup,
) {
	defer wg.Done()

	ticker := time.NewTicker(config.gossip_heartbeat_interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return

		case <-ticker.C:
			// Increment heartbeat for self
			m.mutex.Lock()

			self_id := m.SelfID
			info := m.MembersList[self_id]

			info.HeartBeat++
			m.MembersList[self_id] = info

			m.mutex.Unlock()

			current_members := m.get_current_members()

			selected_peers := select_peers(
				current_members,
				config.gossip_fanout,
				m.get_self_id(),
			)

			gossip_msg := GossipMessage{
				Type:       "gossip",
				Senders_id: m.get_self_id(),
				Members:    current_members,
			}

			for _, peer := range selected_peers {
				peer_addr := &net.UDPAddr{
					IP:   net.ParseIP(peer.IP),
					Port: peer.Port,
				}

				if err := send_gossip_message(socket, peer_addr, gossip_msg); err != nil {
					logger.Warn(
						"Error sending gossip",
						"peer", peer,
						"error", err,
					)
					continue
				}
			}
		}
	}
}

// Merge incoming gossip into local membership.
// Returns true if our own process session had to be replaced.
func (m *MembershipService) handle_gossip(
	incoming_members map[Member_ID]MemberInfo,
	logger *slog.Logger,
) bool {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	now := time.Now()
	restarted := false

	for id, incoming := range incoming_members {
		if incoming.State != "alive" && incoming.State != "down" {
			continue
		}

		// Handle information about our own machine
		if id.IP == m.SelfID.IP && id.Port == m.SelfID.Port {

			// Ignore reports about older local sessions
			if id != m.SelfID {
				continue
			}

			// Someone marked our current session DOWN.
			// Create a NEW session timestamp to refute it.
			if incoming.State == "down" {
				old_id := m.SelfID
				m.deadSessions[old_id] = struct{}{}
				delete(m.MembersList, old_id)

				new_time := now.UTC().Round(0)
				if !new_time.After(old_id.Time) {
					new_time = old_id.Time.Add(time.Nanosecond)
				}

				new_id := Member_ID{
					IP:   old_id.IP,
					Port: old_id.Port,
					Time: new_time,
				}

				m.SelfID = new_id
				m.MembersList[new_id] = MemberInfo{
					HeartBeat:       0,
					State:           "alive",
					local_last_seen: &now,
				}

				restarted = true
				logger.Info("Refuted DOWN report with new session")
			}

			continue
		}

		// Never resurrect an already dead process session
		if _, dead := m.deadSessions[id]; dead {
			continue
		}

		current, exists := m.MembersList[id]

		// New session discovered through gossip
		if !exists {
			m.add_new_member(id, incoming, now)
			continue
		}

		// Same session: DOWN cannot be overridden by ALIVE
		if current.State == "down" {
			continue
		}

		// Incoming DOWN report
		if incoming.State == "down" {
			if incoming.HeartBeat >= current.HeartBeat {
				current.State = "down"
				current.local_failed_since = &now

				m.MembersList[id] = current
				m.deadSessions[id] = struct{}{}

				logger.Info("Member marked DOWN through gossip", "member", id)
			}
			continue
		}

		// Same session + ALIVE: accept newer heartbeat
		if incoming.HeartBeat > current.HeartBeat {
			current.HeartBeat = incoming.HeartBeat
			current.local_last_seen = &now

			m.MembersList[id] = current
		}
	}

	return restarted
}

func (m *MembershipService) handle_failure_detection(
	ctx context.Context,
	socket *net.UDPConn,
	logger *slog.Logger,
	config Config,
	wg *sync.WaitGroup,
) {
	defer wg.Done()

	ticker := time.NewTicker(500 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return

		case <-ticker.C:
			m.mutex.Lock()
			now := time.Now()

			for id, info := range m.MembersList {
				// Don't check self or already DOWN nodes
				if id == m.SelfID || info.State == "down" {
					continue
				}

				if info.local_last_seen == nil {
					info.local_last_seen = &now
					m.MembersList[id] = info
					continue
				}

				// No heartbeat advancement beyond timeout
				if now.Sub(*info.local_last_seen) > config.gossip_failure_timeout {
					info.State = "down"
					info.local_failed_since = &now

					m.MembersList[id] = info

					// Remember this specific process session
					m.deadSessions[id] = struct{}{}

					logger.Warn("Member marked DOWN", "member", id)
				}
			}

			m.mutex.Unlock()
		}
	}
}

func (m *MembershipService) print_membership(logger *slog.Logger) {
	ticker := time.NewTicker(150 * time.Millisecond)
	defer ticker.Stop()

	for range ticker.C {
		current_members := m.get_current_members()
		logger.Info("Current Membership List:")
		for id, info := range current_members {
			logger.Info("Member", "id", id, "state", info.State)
		}
	}
}

func start_membership(
	ctx context.Context,
	config Config,
	logger *slog.Logger,
) {
	logger.Info("Starting Membership")

	if config.my_ip == nil ||
		config.gossip_heartbeat_interval <= 0 ||
		config.gossip_failure_timeout <= 0 ||
		config.join_retry_interval <= 0 {
		logger.Error("Invalid membership configuration")
		return
	}

	var membership_service MembershipService

	membership_service.MembersList = make(map[Member_ID]MemberInfo)
	membership_service.deadSessions = make(map[Member_ID]struct{})

	membership_service.member_init(
		config.my_ip.String(),
		config.port,
	)

	socket, err := net.ListenUDP("udp", &net.UDPAddr{
		IP:   *config.my_ip,
		Port: config.port,
	})

	if err != nil {
		logger.Error("Failed to create UDP listener", "error", err)
		return
	}

	defer socket.Close()

	logger.Info(
		"Membership Service Listening on",
		"IP", config.my_ip.String(),
		"Port", config.port,
	)

	var wg sync.WaitGroup
	wg.Add(3)

	go membership_service.handle_gossip_receive(
		ctx, socket, logger, config, &wg,
	)

	go membership_service.handle_gossip_send(
		ctx, socket, logger, config, &wg,
	)

	go membership_service.handle_failure_detection(
		ctx, socket, logger, config, &wg,
	)

	// Wait for cancellation
	<-ctx.Done()

	// Close socket to unblock receiver
	socket.Close()

	// Wait for all three goroutines to finish
	wg.Wait()

	logger.Info("Membership Service Stopped")
}