import socket
import threading
import time
import copy
from typing import Dict, Any, List

# Scapy is a third-party library. Install it with: pip install scapy
try:
    from scapy.all import sniff, get_if_list
    from scapy.config import conf
except ImportError:
    print("Scapy is not installed. Please run 'pip install scapy'")
    exit(1)

# Shared statistics dictionary, protected by a lock
# This is the Python equivalent of your 'volatile struct ss_t stat'
stats: Dict[str, Dict[str, Any]] = {
    "pcap": {
        "name": "PCAP",
        "counter": 0,
        "period_us": 0,
        "min_us": 0,
        "max_us": 0,
    },
    "socket": {
        "name": "SOCKET",
        "counter": 0,
        "period_us": 0,
        "min_us": 0,
        "max_us": 0,
    },
}
stats_lock = threading.Lock()

# --- Configuration ---
#TARGET_IP = "192.168.2.2"
TARGET_IP = "0.0.0.0"
TARGET_PORT = 60012
BPF_FILTER = f"udp and port {TARGET_PORT}"
STATS_UPDATE_INTERVAL_S = 0.5
WARMUP_PACKETS = 100 # Number of packets to ignore at the beginning

def list_interfaces() -> List[Any]:
    """Lists available network interfaces in a user-friendly way."""
    print("Available network interfaces:")
    # We use get_if_list() to get the names and then access full details
    # from scapy's configuration object for better descriptions.
    interfaces = []
    if_list = get_if_list()
    for i, if_name in enumerate(if_list):
        iface = conf.ifaces.get(if_name)
        if iface:
            description = iface.description or "No description"
            print(f"{i + 1}. {if_name} ({description})")
            interfaces.append(iface)
    return interfaces

def pcap_capture_thread(interface_name: str) -> None:
    """
    Captures packets using Scapy and calculates inter-packet latency.
    This function replaces the C++ CaptureThreadProc and packet_handler.
    """
    print(f"PCAP thread started on interface '{interface_name}'...")
    last_packet_time: float = 0.0

    def process_packet(packet) -> None:
        """Callback function for each captured packet."""
        nonlocal last_packet_time
        
        current_time = packet.time
        
        # We use a nonlocal variable to store the timestamp of the previous packet
        if last_packet_time == 0.0:
            last_packet_time = current_time
            return

        delta_s = current_time - last_packet_time
        delta_us = int(delta_s * 1_000_000)
        last_packet_time = current_time

        with stats_lock:
            stats["pcap"]["counter"] += 1
            
            # Start calculating stats only after a warm-up period
            if stats["pcap"]["counter"] > WARMUP_PACKETS:
                stats["pcap"]["period_us"] = delta_us
                
                current_min = stats["pcap"]["min_us"]
                current_max = stats["pcap"]["max_us"]
                
                if current_min == 0 or delta_us < current_min:
                    stats["pcap"]["min_us"] = delta_us
                if delta_us > current_max:
                    stats["pcap"]["max_us"] = delta_us
            else:
                # Initialize min/max during warm-up
                stats["pcap"]["min_us"] = delta_us
                stats["pcap"]["max_us"] = delta_us


    # sniff is the Scapy equivalent of pcap_loop.
    # store=0 ensures packets are not stored in memory, for performance.
    sniff(iface=interface_name, filter=BPF_FILTER, prn=process_packet, store=0)


def socket_receive_thread(host_ip: str, port: int) -> None:
    """
    Receives UDP packets on a standard socket and measures latency.
    This function replaces the C++ RcvThreadProc.
    """
    print(f"Socket thread started on {host_ip}:{port}...")
    
    # time.perf_counter() is the Python equivalent of QueryPerformanceCounter
    last_receive_time: float = 0.0
    
    try:
        # AF_INET for IPv4, SOCK_DGRAM for UDP
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.bind((host_ip, port))
    except socket.error as e:
        print(f"Error binding socket: {e}")
        return

    while True:
        try:
            # Buffer size, similar to the C++ code
            sock.recvfrom(2048)
            current_time = time.perf_counter()

            if last_receive_time == 0.0:
                last_receive_time = current_time
                continue

            delta_s = current_time - last_receive_time
            delta_us = int(delta_s * 1_000_000)
            last_receive_time = current_time

            with stats_lock:
                stats["socket"]["counter"] += 1
                
                # Start calculating stats only after a warm-up period
                if stats["socket"]["counter"] > WARMUP_PACKETS:
                    stats["socket"]["period_us"] = delta_us

                    current_min = stats["socket"]["min_us"]
                    current_max = stats["socket"]["max_us"]
                    
                    if current_min == 0 or delta_us < current_min:
                        stats["socket"]["min_us"] = delta_us
                    if delta_us > current_max:
                        stats["socket"]["max_us"] = delta_us
                else:
                    # Initialize min/max during warm-up
                    stats["socket"]["min_us"] = delta_us
                    stats["socket"]["max_us"] = delta_us

        except socket.error as e:
            # This can happen if the socket is closed or another error occurs
            print(f"\nSocket error: {e}")
            with stats_lock:
                # Store the error code, similar to the WSAGetLastError logic
                stats["socket"]["counter"] = f"Error: {e.errno}"
            break
            
    sock.close()


def main():
    """Main function to set up threads and display statistics."""
    
    interfaces = list_interfaces()
    if not interfaces:
        print("\nNo interfaces found! Make sure Scapy is installed and you have privileges.")
        return

    try:
        choice = int(input(f"Enter the interface number (1-{len(interfaces)}): "))
        if not 1 <= choice <= len(interfaces):
            print("Invalid number.")
            return
        selected_interface = conf.ifaces.get(get_if_list()[choice - 1])
    except (ValueError, IndexError):
        print("Invalid input.")
        return

    print(f"\nListening on '{selected_interface.description}'...")

    # Create and start the PCAP capture thread
    pcap_thread = threading.Thread(
        target=pcap_capture_thread,
        args=(selected_interface.name,),
        daemon=True  # Threads will exit when the main program exits
    )

    # Create and start the socket receive thread
    socket_thread = threading.Thread(
        target=socket_receive_thread,
        args=(TARGET_IP, TARGET_PORT),
        daemon=True
    )
    
    pcap_thread.start()
    socket_thread.start()
    
    print("\nStarting measurements... Press Ctrl+C to stop.\n")
    time.sleep(2) # Give threads time to initialize

    try:
        while True:
            # Use deepcopy to get a snapshot of the stats, preventing race
            # conditions during printing. This is the Pythonic way to do
            # what you did with memcpy.
            with stats_lock:
                stats_snapshot = copy.deepcopy(stats)

            # Format the output string
            line = "\r"
            for key in ["pcap", "socket"]:
                s = stats_snapshot[key]
                line += (
                    f"{s['name']}: {str(s['counter']):>6} | "
                    f"period={s['period_us']:>7}µs | "
                    f"min={s['min_us']:>7}µs | "
                    f"max={s['max_us']:>7}µs    "
                )
            
            print(line, end="", flush=True)
            time.sleep(STATS_UPDATE_INTERVAL_S)

    except KeyboardInterrupt:
        print("\n\nStopping analyzer.")


if __name__ == "__main__":
    main()