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)