import sys
import time
import threading
from scapy.all import Ether, IP, UDP, send, sendp, Raw, sniff

global initial_timestamp
initial_timestamp = int(time.time() * 1000)

# Simulated ACK sending function
def send_ack(destination_IP, my_port, sender_port):
    ack_data = ""
    ###################
    # With localhost:
    ack_packet = IP(dst=destination_IP) / UDP(dport=sender_port, sport=my_port) / Raw(load=str(ack_data))
    send(ack_packet, verbose=False)
    ###################
    # With netns:
    # ack_packet = Ether(dst="00:00:00:00:00:04", src="00:00:00:00:00:01") / IP(dst="127.0.0.1") / UDP(dport=5551) / Raw(load=str(ack_data))
    # sendp(ack_packet, verbose=False)
    ###################
    print(f"Sent ACK: {ack_data}")

# Process incoming packets
def process_packet(pkt):
    ################### Save timestamps ###################
    if Raw in pkt:
        with open(save_file, "a") as f: # open as "a" for append
            global initial_timestamp
            receival_time = str(int(time.time() * 1000) - initial_timestamp) # time of arrival at receiver
            send_time = pkt[Raw].load # time it was sent from the sender
            ################### Conversion & error handling ###################
            try:
                send_time = int(send_time.decode("utf-8").split("X")[0]) # decode the payload and stop at "X"
            except ValueError: # For the last packet without a timestamp but "END" as payload
                print(f"END for {my_port}")
                send_ack(destination_IP, my_port, sender_port)
                return
            except Exception as e:
                print(f"Decoding error at {my_port}")
                send_ack(destination_IP, my_port, sender_port)
                return
            ################### Write data ###################
            aoi = int(receival_time) - send_time
            f.write(f"{send_time},{receival_time},{aoi}\n")

    send_ack(destination_IP, my_port, sender_port)

# Main loop
def main(destination_IP, my_port, sender_port):
    ###################
    # With localhost:
    sniff(iface="lo", filter=f"udp and dst port {my_port}", prn=process_packet, store=False)
    ###################
    # With netns:
    # sniff(iface="ns-g", filter=f"udp and dst port {my_port}", prn=process_packet, store=False)
    ###################

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: python receive_and_ack.py <destination_IP> <my_port> <sender_port>")
        sys.exit(1)

    destination_IP = sys.argv[1]                     # Clear
    my_port = int(sys.argv[2])                       # Port to listen on
    sender_port = int(sys.argv[3])                   # Port to send acks to
    save_file = f"timestamps_{my_port}"
    main(destination_IP, my_port, sender_port)