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"
IP = socket.gethostbyname(socket.gethostname())
SERVER_PORT = 6767
INTRODUCER_FLAG = (INTRODUCER_IP == IP)
HEART_BEAT_INTERVAL = 5
GOSSIP_INTERVAL = 10.0
FAILURE_TIMEOUT = 3.0

members_list = {}
members_list_lock = threading.Lock()

SELF_ID = (IP, SERVER_PORT, utc_time())
members_list[SELF_ID] = {
    "heart_beat": 1,
    "last_seen": utc_time()
}
print(f"Starting peer: IP={IP}, port={SERVER_PORT}, introducer={INTRODUCER_FLAG}, self_id={SELF_ID}")

stop_event = threading.Event()

server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
server_socket.bind((IP, SERVER_PORT))
server_socket.settimeout(1.0)

def worker_1():
    if not INTRODUCER_FLAG:
        join_payload = {
            "type": "JOIN",
            "sender_id": SELF_ID
        }
        server_socket.sendto(json.dumps(join_payload).encode('utf-8'), (INTRODUCER_IP, SERVER_PORT))
    
    while not stop_event.is_set():
        try:
            data, addr = server_socket.recvfrom(65535) 
            
            try:
                msg = json.loads(data.decode('utf-8'))
            except json.JSONDecodeError:
                continue
                
            msg_type = msg.get("type")
            print(f"Processing message type={msg_type} from {addr}")
            
            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 join_id not in members_list:
                        members_list[join_id] = {
                            "heart_beat": 1,
                            "last_seen": utc_time()
                        }
                        print(f"Added member: {join_id}")
                    
                    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"]
                        })
                        
                print(f"New member joined: {join_id}")
                
                # Send the initial state directly 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))
                print(f"Sent initial membership state to {(join_ip, join_port)}: {reply_payload}")

            elif msg_type == "GOSSIP":
                members_json = msg.get('members', [])

                with members_list_lock:
                    for member in members_json:
                        member_id = (member['ip'], member['port'], member['timestamp'])
                        member_hb = member['heart_beat']

                        if member_id == SELF_ID:
                            continue

                        if member_id not in members_list:
                            members_list[member_id] = {
                                "heart_beat": member_hb,
                                "last_seen": utc_time()
                            }
                            print(f"Added gossip member: {member_id}, heartbeat={member_hb}")
                        elif member_hb > members_list[member_id]["heart_beat"]:
                            members_list[member_id]["heart_beat"] = member_hb
                            members_list[member_id]["last_seen"] = utc_time()
                            print(f"Updated member: {member_id}, heartbeat={member_hb}")
                
        except socket.timeout:
            continue

def worker_2():
    while not stop_event.is_set():
        with members_list_lock:
            members_list[SELF_ID]["heart_beat"] += 1
            print(f"Updated my heartbeat: {members_list[SELF_ID]['heart_beat']}")
        time.sleep(HEART_BEAT_INTERVAL)

def worker_3():
    while not stop_event.is_set():
        with members_list_lock:
            peers = list(members_list.keys())
            
            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"]
                })
            
        if SELF_ID in peers:
            peers.remove(SELF_ID)
            
        if peers:
            random_peer = random.choice(peers)
            peer_ip, peer_port, _ = random_peer
            
            payload = {
                "type": "GOSSIP",
                "sender_id": SELF_ID,
                "members": members_payload
            }
            
            json_string = json.dumps(payload)
            print(f"Sending GOSSIP to {(peer_ip, peer_port)}: {payload}")
            server_socket.sendto(json_string.encode('utf-8'), (peer_ip, peer_port))
            
        time.sleep(GOSSIP_INTERVAL)

#only for debuging
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)
#only for debugging
t4 = threading.Thread(target=worker_monitoring)

t1.start()
t2.start()
t3.start()
t4.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()
server_socket.close()