import logging
import os
from bcc import BPF

logger = logging.getLogger(__name__)

script_dir = os.path.dirname(os.path.abspath(__file__))


def int_to_ip(ip):
    """Convert an int IP address to dotted-quad format (e.g, 192.168.1.0)."""
    return f"{ip & 0xFF}.{ip >> 8 & 0xFF}.{ip >> 16 & 0xFF}.{ip >> 24 & 0xFF}"


with open(os.path.join(script_dir, "kprobe.c"), 'r') as f:
    """Read ebpf C program."""
    bpf_program = f.read()

def load_bpf_program():
    """Load ebpf program."""
    b = BPF(text=bpf_program)
    b.attach_kprobe(event="tcp_retransmit_skb", fn_name="count_tcp_retransmit")
    logger.info("Probe loaded: counting TCP retransmissions per flow... ")
    return b

def probe(b, func_record, timeout=1000):
    """Pool perf buffer for data."""
    rtx_count = b["rtx_count"]
    logger.info("Probe called. Number of records: %d", len(rtx_count.items()))
    for flow, count in rtx_count.items():
        func_record(int_to_ip(flow.src_ip),
                    int_to_ip(flow.dst_ip),
                    flow.src_port,
                    flow.dst_port,
                    count.value)

    rtx_count.clear()
    b.perf_buffer_poll(timeout=timeout)