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

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


# Parameters
Q_max = 10  # Maximum queue size
age_threshold = 1.0  # Delta_T in seconds

# Sender's state
last_update_time = None  # Last time an ACK was received
N = 1  # Active senders (received from ACK)
queue_utilization = 0  # Received as Q_i/Q_max
delta_c = 0  # Time since the last ACK (calculated using T_ACK)


# Function to calculate P_s
def calculate_probability(N, queue_utilization, delta_c):
    if queue_utilization < 1:
        return 1.0  # Full probability if the queue is underutilized
    elif queue_utilization == 1:
        # Calculate f(Delta_i)
        f_delta = max((delta_c - age_threshold) / age_threshold, 0)
        return (Q_max / N) + f_delta  # Combine probabilities


# Function to send a packet
def send_packet(target_ip, receiver_port, my_port, dscp_value, packet_size, FIN_flag):
    ether_layer = Ether(dst="00:00:00:00:00:04", src="00:00:00:00:00:01")
    ip_layer = IP(dst=target_ip, tos=dscp_value << 2)
    udp_layer = UDP(dport=receiver_port, sport=my_port)

    # # Socket to listen for the ACKs
    # ack_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    # ack_socket.settimeout(2)                                            # 2 seconds timer to receive the ack. Otherwise considered lost
    # ack_socket.bind(("", my_port))

    # timestamp = str(int(time.time()))                                   # Convert to a string for embedding

    global initial_timestamp
    timestamp = str(int(time.time() * 1000) - initial_timestamp)

    if not FIN_flag:
        payload = timestamp + "X" * max(
            (packet_size - len(timestamp)), 0
        )  # max to prevent possible errors when choosing too small packet size
    else:
        payload = "END" + "X" * max((packet_size - len(timestamp)), 0)

    # packet = ether_layer / ip_layer / udp_layer / Raw(load=payload)
    packet = ip_layer / udp_layer / Raw(load=payload)
    # sendp(packet, verbose=False)
    send(packet, verbose=False)
    print(".", end="", flush=True)


# Listen for ACKs
def process_ack(pkt):
    global last_update_time, N, queue_utilization, delta_c

    if Raw in pkt:
        ack_data = pkt[Raw].load.decode("utf-8")
        N, queue_utilization, T_ACK = eval(ack_data)
        last_update_time = T_ACK
        print(f"Received ACK: N={N}, Utilization={queue_utilization}, T_ACK={T_ACK}")


# Main loop
def main(target_ip, receiver_port, my_port, dscp_value, nb_packets, packet_size):
    global delta_c
    packets_sent = 0

    while packets_sent < nb_packets:
        # Update delta_c (time since the last ACK)
        if last_update_time:
            delta_c = time.time() - last_update_time
        else:
            delta_c = float("inf")  # No ACK received yet

        # Calculate sending probability
        P_s = calculate_probability(N, queue_utilization, delta_c)

        # Decide whether to send a packet
        if random.random() < P_s:
            FIN_flag = False
            if packets_sent == nb_packets - 1:
                FIN_flag = True

            send_packet(
                target_ip, receiver_port, my_port, dscp_value, packet_size, FIN_flag
            )
            packets_sent += 1
            print(f"Packet sent with probability P_s={P_s:.2f}")
        else:
            print(f"Packet skipped with probability P_s={P_s:.2f}")

        # Wait for a short interval before the next attempt
        time.sleep(0.1)

    print(f"Finished sending {nb_packets} packets to {target_ip}:{receiver_port}")


if __name__ == "__main__":
    if len(sys.argv) != 6:
        print(
            "Usage: python send_dscp_packets.py <target_ip> <my_port> <receiver_port> <nb_packets> <packet_size(in Bytes)>"
        )
        sys.exit(1)

    target_ip = sys.argv[1]  # IP of the receiver
    my_port = int(sys.argv[2])  # Sender port
    receiver_port = int(sys.argv[3])  # Receiver port
    nb_packets = int(sys.argv[4])  # Total number of packets to send
    packet_size = int(sys.argv[5]) - 42  # Packet size in Bytes - 42 (previous headers)
    dscp_value = 0  # Hardcoded DSCP value for now

    # main(target_ip, receiver_port, my_port, dscp_value, nb_packets)

    threading.Thread(
        target=main,
        args=(target_ip, receiver_port, my_port, dscp_value, nb_packets, packet_size),
        daemon=True,
    ).start()
    ###################
    # With localhost:
    sniff(
        iface="lo", filter=f"udp and dst port {my_port}", prn=process_ack, store=False
    )  # , timeout=0.1)
    ###################
# With netns:
# sniff(iface="ns-r", filter=f"udp and dst port {my_port}", prn=process_ack, store=False)
###################

""" Archive:

# Add padding:
    payload = timestamp + "X" * (packet_size - len(timestamp))




"""