package main

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

// 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 {
	heart_beat          int
	state               string
	incarnation_counter int
	local_failed_since  *time.Time
	local_last_seen     *time.Time
}

type MembershipService struct {
	members_list map[Member_ID]MemberInfo
	mutex        sync.RWMutex
}

type GossipMessageType int


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

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.members_list {
		current_members[id] = info
	}
	return current_members
}

func get_all_member_ids(members map[Member_ID]MemberInfo) []Member_ID {
	member_ids := make([]Member_ID, 0, len(members) - 1) // Exclude self
	for id := range members {
		
	}
	return member_ids
}

func (m *MembershipService) member_init(ip string, port int) {
	my_id := Member_ID{ip, port, time.Now()}
	members_info := MemberInfo{0, "alive", 0, nil, nil}
	m.mutex.Lock()
	defer m.mutex.Unlock()
	m.members_list[my_id] = members_info
}

func (m *MembershipService) handle_join(new_member Member_ID, logger *slog.Logger, socket *net.UDPConn) {

	m.mutex.Lock()
	m.members_list[new_member] = MemberInfo{0, "alive", 0, nil, nil}
	m.mutex.Unlock()

	//ack payload with teh current membership list
	gossip_msg := GossipMessage{
		Type:    "gossip",
		Senders_id: Member_ID{},
		Members: m.get_current_members(),
	}
	data, err := json.Marshal(gossip_msg)
	if err != nil {
		logger.Error("Error marshaling gossip message", "error", err)
		return
	}

	_, err = socket.WriteToUDP(data, &net.UDPAddr{IP: net.ParseIP(new_member.ip), Port: new_member.port})
	if err != nil {
		logger.Error("Error sending gossip message", "error", err)
		return
	}
}


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)
	//send the join message before processing the gossip messages
	joined := (config.introducer_ip == config.my_ip.String()) // Introducer node doesn't need to join
	last_join_time := time.Time{}

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

		if !joined && time.Since(last_join_time) >= config.join_retry_interval {
			// Send join message to introducer
			join_msg := GossipMessage{
				Type:    "join",
				Senders_id: Member_ID{ip: config.my_ip.String(), port: config.port, time: time.Now()},
				Members: map[Member_ID]MemberInfo{},
			}
			data, err := json.Marshal(join_msg)
			if err != nil {
				logger.Error("Error marshaling join message", "error", err)
				return
			}
			introducer_addr := &net.UDPAddr{
				IP:   net.ParseIP(config.introducer_ip),
				Port: config.port,
			}
			_, err = socket.WriteToUDP(data, introducer_addr)
			if err != nil {
				logger.Error("Error sending join message", "error", err)
				return
			}
			last_join_time = time.Now()
		}

		socket.SetReadDeadline(time.Now().Add(1 * time.Second))
		n, _, err := socket.ReadFromUDP(buffer)
		if err != nil {
			if ctx.Err() != nil {
				return
			}
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				// Timeout occurred, continue to the next iteration to check for context cancellation
				continue
			}
			logger.Error("Error reading from UDP socket", "error", err)
			return
		}
		
		var gossip_msg GossipMessage
		err = json.Unmarshal(buffer[:n], &gossip_msg)
		if err != nil {
			logger.Error("Error unmarshaling gossip message", "error", err)
			return
		}

		switch gossip_msg.Type {
		case "gossip":
			m.handle_gossip(gossip_msg.Members, logger)
		case "join":
			if config.introducer_ip == config.my_ip.String() {
				m.handle_join(gossip_msg.Senders_id, logger, socket)
				joined = true
			} else {
				logger.Error("Received join message but not the introducer")
			}
		default:
			logger.Error("Unknown message type received", "type", gossip_msg.Type)
		}
	}
}

func select_peers(current_members map[Member_ID]MemberInfo, fanout int) []Member_ID {
	if len(current_members) <= fanout {
		return get_all_member_ids(current_members)
	}
	
	member_ids := get_all_member_ids(current_members)
	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 heart beat for self
			m.mutex.Lock()
			for id, info := range m.members_list {
				if id.ip == config.my_ip.String() && id.port == config.port {
					info.heart_beat++
					m.members_list[id] = info
					break
				}
			}
			m.mutex.Unlock()

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

			for _, peer := range selected_peers {
				gossip_msg := GossipMessage{
					Type:    "gossip",
					Members: current_members,
				}
				data, err := json.Marshal(gossip_msg)
				if err != nil {
					logger.Error("Error marshaling gossip message", "error", err)
					return
				}
				// Send the gossip message to the selected peer
				peer_addr := &net.UDPAddr{
					IP:   net.ParseIP(peer.ip),
					Port: peer.port,
				}
				_, err = socket.WriteToUDP(data, peer_addr)
				if err != nil {
					logger.Error("Error sending gossip message", "error", err)
					return
				}
			}
		}
	}
}

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:
			// TODO: Implement failure detection logic
		}
	}
}

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

	var membership_service MembershipService
	membership_service.members_list = make(map[Member_ID]MemberInfo)
	membership_service.member_init(config.my_ip.String(), config.port)

	// 2. Check the error on socket creation
	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
	}

	// 5. Keep defer socket.Close() since the function remains blocking
	defer socket.Close()

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

	// 6. Goroutine tracking with WaitGroup
	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)

	wg.Wait()
}