from datetime import datetime, timezone
import socket
import threading
import time
import random
import json 

def utc_time():
    return datetime.now(timezone.utc).timestamp()

INTRODUCER_IP = "172.22.156.3" #hardcoded ip
IP = socket.gethostbyname(socket.gethostname()) #the vm IP address
SERVER_PORT = 8989
INTRODUCER_FLAG = (INTRODUCER_IP == IP)
HEART_BEAT_INTERVAL = 0.25
GOSSIP_INTERVAL = 0.25
FAILURE_TIMEOUT = 2.5
FAILURE_T_CLEANUP = 10.0
RETRY_INTERVAL = 5.0
FANOUT = 1

#message json example:
# {
#     type: "GOSSIP", or "JOIN", ...
#     sender_id: (ip, port, timestamp),
#     members: [
#         {
#             ip: "
#             port: 8989,
#             timestamp: 1234567890.123456,
#             heart_beat: 1,
#             state: "alive"
#         },
#         ...

members_list = {}
dead_members = set() #[BASICALLY VERSION TRACKING]this tracks members that have been marked as failed and removed from the membership list.
members_list_lock = threading.Lock()

#adding its self to the membership list
SELF_ID = (IP, SERVER_PORT, utc_time())
members_list[SELF_ID] = {
    "heart_beat": 1,
    "local_last_seen": utc_time(),
    "state": "alive",
    "local_failed_since" : None,
    "local_suspected_since" : None
}

stop_event = threading.Event() #trigger to stop all threads when the program is exiting

#one socket for all threads to use
server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
server_socket.bind((IP, SERVER_PORT))
server_socket.settimeout(1.0)

#responsible for handling incoming gossip messages, join requests on the introducer, and updating the membership list based on the received messages
def worker_1():
    has_joined = INTRODUCER_FLAG #if this is the introducer, it has already joined
    last_join_time = 0 #flag to track the last time a join request was sent to retry
    
    while not stop_event.is_set():
        # If the node has not joined and enough time has passed since the last join attempt, send a JOIN request to the introducer
        if not has_joined and (utc_time() - last_join_time > RETRY_INTERVAL):
            join_payload = {
                "type": "JOIN",
                "sender_id": SELF_ID
            }
            server_socket.sendto(json.dumps(join_payload).encode('utf-8'), (INTRODUCER_IP, SERVER_PORT))
            last_join_time = utc_time()

        try:
            data, addr = server_socket.recvfrom(65535) 
            
            msg = json.loads(data.decode('utf-8'))
                
            msg_type = msg.get("type")
            
            #handle join type
            if msg_type == "JOIN":
                join_ip, join_port, join_timestamp = msg["sender_id"]
                join_id = (join_ip, join_port, join_timestamp)
                
                with members_list_lock:
                    #if in dead members, ignore the join request
                    if join_id in dead_members:
                        continue

                    if join_id not in members_list:
                        members_list[join_id] = {
                            "heart_beat": 1,
                            "local_last_seen": utc_time(),
                            "state": "alive",
                            "local_failed_since" : None,
                            "local_suspected_since" : None
                        }
                    
                    members_payload = []
                    for m_id, m_data in members_list.items():
                        members_payload.append({
                            "ip": m_id[0],
                            "port": m_id[1],
                            "timestamp": m_id[2],
                            "heart_beat": m_data["heart_beat"],
                            "state": m_data["state"]
                        })
                
                #send the current membership list to the joining node
                reply_payload = {
                    "type": "GOSSIP",
                    "sender_id": SELF_ID,
                    "members": members_payload
                }
                server_socket.sendto(json.dumps(reply_payload).encode('utf-8'), (join_ip, join_port))

            #handle gossip type
            elif msg_type == "GOSSIP":
                #first time receiving a gossip message, mark the node as joined
                if not has_joined:
                    has_joined = True

                members_json = msg.get('members', [])

                with members_list_lock:
                    for member in members_json:
                        member_id = (member['ip'], member['port'], member['timestamp'])
                        
                        if member_id == SELF_ID or member_id in dead_members:
                            continue

                        incoming_state = member['state']
                        incoming_hb = member['heart_beat']
                            
                        local_member = members_list[member_id]

                        if local_member is None:
                            if incoming_state == "failed":
                                dead_members.add(member_id)
                            else:
                                members_list[member_id] = {
                                    "heart_beat": incoming_hb,
                                    "local_last_seen": utc_time(),
                                    "state": incoming_state,
                                    "local_failed_since" : None,
                                    "local_suspected_since" : None
                                }
                            continue

                        if local_member["state"] == "failed":
                            continue

                        if incoming_state == "failed":
                            local_member["state"] = "failed"
                            local_member["local_failed_since"] = utc_time()
                            print(f"Member {member_id} marked as failed due to gossip.")
                        elif incoming_hb > local_member["heart_beat"]:
                            local_member["heart_beat"] = incoming_hb
                            local_member["local_last_seen"] = utc_time()
                
        except socket.timeout:
            continue

#handles incrementing the local node's heartbeat counter at regular intervals
def worker_2():
    while not stop_event.is_set():
        with members_list_lock:
            members_list[SELF_ID]["heart_beat"] += 1
        time.sleep(HEART_BEAT_INTERVAL)

#handles sending gossip messages to a subset of alive peers at regular intervals
def worker_3():
    peer_queue = []
    
    while not stop_event.is_set():
        with members_list_lock:
            current_peers = [member_id for member_id in members_list.keys() if member_id != SELF_ID and members_list[member_id]["state"] == "alive"]
            
            members_payload = []
            for m_id, m_data in members_list.items():
                members_payload.append({
                    "ip": m_id[0],
                    "port": m_id[1],
                    "timestamp": m_id[2],
                    "heart_beat": m_data["heart_beat"],
                    "state": m_data["state"]
                })
        
        if current_peers:
            #remove any peers from the queue that are no longer alive 
            peer_queue = [p for p in peer_queue if p in current_peers]
            
            selected_peers = []
            num_to_select = min(FANOUT, len(current_peers))
            
            for _ in range(num_to_select):
                if not peer_queue:
                    peer_queue = list(current_peers)
                    random.shuffle(peer_queue)
                
                if peer_queue:
                    selected_peers.append(peer_queue.pop(0))
            
            for peer in selected_peers:
                peer_ip, peer_port, _ = peer
                
                payload = {
                    "type": "GOSSIP",
                    "sender_id": SELF_ID,
                    "members": members_payload
                }
                
                json_string = json.dumps(payload)
                server_socket.sendto(json_string.encode('utf-8'), (peer_ip, peer_port))
            
        time.sleep(GOSSIP_INTERVAL)

# handle failure detection and cleanup of failed members from the membership list
def worker_4():
    while not stop_event.is_set():
        current_time = utc_time()
        to_remove = []
        
        with members_list_lock:
            for member_id, member_data in members_list.items():
                if member_id == SELF_ID:
                    continue
                
                if member_data["state"] == "alive":
                    last_seen = member_data["local_last_seen"]
                    time_since_last_seen = current_time - last_seen
                    
                    if time_since_last_seen > FAILURE_TIMEOUT:
                        member_data["state"] = "failed"
                        member_data["local_failed_since"] = current_time
                
                last_failed_since = member_data.get("local_failed_since")
                if last_failed_since is not None and (current_time - last_failed_since) > FAILURE_T_CLEANUP:
                    to_remove.append(member_id)
            
            for member_id in to_remove:
                dead_members.add(member_id)
                del members_list[member_id]
                
        time.sleep(0.5)

#for monitoring purposes
def worker_monitoring():
    while not stop_event.is_set():
        with members_list_lock:
            local_ml = {
                member_id: member_data.copy()
                for member_id, member_data in members_list.items()
            }
        
        with open("membership_list.txt", "w") as f:
            for member_id, member_data in local_ml.items():
                f.write(f"{member_id}: {member_data}\n")
        time.sleep(0.5)

t1 = threading.Thread(target=worker_1)
t2 = threading.Thread(target=worker_2)
t3 = threading.Thread(target=worker_3)
t4 = threading.Thread(target=worker_4)
t5 = threading.Thread(target=worker_monitoring)

t1.start()
t2.start()
t3.start()
t4.start()
t5.start()

try:
    while not stop_event.is_set():
        time.sleep(0.5)
except KeyboardInterrupt:
    print("\nExiting...")
    stop_event.set()

t1.join()
t2.join()
t3.join()
t4.join()
t5.join()
server_socket.close()