#!/usr/bin/env python3
"""eBPF-based UDP tracer for MP-QUIC performance measurement.

Traces kernel-level UDP send/receive latency and packet drops.
Requires root privileges and Linux with BCC installed.
"""

from bcc import BPF
import time
import argparse
import signal
import sys

bpf_text = """
#include <uapi/linux/ptrace.h>
#include <net/sock.h>
#include <net/inet_sock.h>
#include <bcc/proto.h>

#define MPQUIC_PORT __TARGET_PORT__

BPF_HASH(send_start, u32, u64);
BPF_HASH(recv_start, u32, u64);

BPF_PERF_OUTPUT(events);

struct event_t {
    u64 timestamp_ns;
    u64 delta_ns;
    u32 pid;
    u16 sport;
    u16 dport;
    u8  event_type;  // 0=SEND, 1=RECV, 2=DROP
};

int trace_udp_sendmsg(struct pt_regs *ctx, struct sock *sk) {
    struct inet_sock *inet = (struct inet_sock *)sk;
    u16 dport = 0;
    bpf_probe_read_kernel(&dport, sizeof(dport), &inet->inet_dport);
    dport = ntohs(dport);

    // Filter: only trace traffic to/from MP-QUIC port
    u16 sport = 0;
    bpf_probe_read_kernel(&sport, sizeof(sport), &inet->inet_sport);
    sport = ntohs(sport);

    if (dport != MPQUIC_PORT && sport != MPQUIC_PORT)
        return 0;

    u32 pid = bpf_get_current_pid_tgid();
    u64 ts = bpf_ktime_get_ns();
    send_start.update(&pid, &ts);
    return 0;
}

int trace_udp_sendmsg_ret(struct pt_regs *ctx) {
    u32 pid = bpf_get_current_pid_tgid();
    u64 *tsp = send_start.lookup(&pid);
    if (tsp == 0)
        return 0;

    u64 delta = bpf_ktime_get_ns() - *tsp;
    send_start.delete(&pid);

    struct event_t event = {};
    event.timestamp_ns = bpf_ktime_get_ns();
    event.delta_ns = delta;
    event.pid = pid;
    event.event_type = 0;  // SEND
    events.perf_submit(ctx, &event, sizeof(event));
    return 0;
}

int trace_udp_recvmsg(struct pt_regs *ctx, struct sock *sk) {
    struct inet_sock *inet = (struct inet_sock *)sk;
    u16 sport = 0;
    bpf_probe_read_kernel(&sport, sizeof(sport), &inet->inet_sport);
    sport = ntohs(sport);

    u16 dport = 0;
    bpf_probe_read_kernel(&dport, sizeof(dport), &inet->inet_dport);
    dport = ntohs(dport);

    if (sport != MPQUIC_PORT && dport != MPQUIC_PORT)
        return 0;

    u32 pid = bpf_get_current_pid_tgid();
    u64 ts = bpf_ktime_get_ns();
    recv_start.update(&pid, &ts);
    return 0;
}

int trace_udp_recvmsg_ret(struct pt_regs *ctx) {
    u32 pid = bpf_get_current_pid_tgid();
    u64 *tsp = recv_start.lookup(&pid);
    if (tsp == 0)
        return 0;

    u64 delta = bpf_ktime_get_ns() - *tsp;
    recv_start.delete(&pid);

    struct event_t event = {};
    event.timestamp_ns = bpf_ktime_get_ns();
    event.delta_ns = delta;
    event.pid = pid;
    event.event_type = 1;  // RECV
    events.perf_submit(ctx, &event, sizeof(event));
    return 0;
}

TRACEPOINT_PROBE(skb, kfree_skb) {
    struct event_t event = {};
    event.timestamp_ns = bpf_ktime_get_ns();
    event.delta_ns = 0;
    event.pid = bpf_get_current_pid_tgid();
    event.event_type = 2;  // DROP
    events.perf_submit(ctx, &event, sizeof(event));
    return 0;
}
"""

EVENT_TYPES = {0: "SEND_LATENCY", 1: "RECV_LATENCY", 2: "DROP"}

# Counters for live summary
stats = {"send_count": 0, "recv_count": 0, "drop_count": 0,
         "send_total_ns": 0, "recv_total_ns": 0}


def main():
    parser = argparse.ArgumentParser(description="eBPF UDP Tracer for MP-QUIC")
    parser.add_argument("--output", type=str, default="measurement.csv",
                        help="CSV output file")
    parser.add_argument("--port", type=int, default=4242,
                        help="MP-QUIC port to filter on")
    args = parser.parse_args()

    program = bpf_text.replace("__TARGET_PORT__", str(args.port))

    print(f"Compiling eBPF program (requires root)...")
    print(f"Filtering on UDP port {args.port}")
    b = BPF(text=program)

    # Attach kprobes for send and receive paths
    b.attach_kprobe(event="udp_sendmsg", fn_name="trace_udp_sendmsg")
    b.attach_kretprobe(event="udp_sendmsg", fn_name="trace_udp_sendmsg_ret")
    b.attach_kprobe(event="udp_recvmsg", fn_name="trace_udp_recvmsg")
    b.attach_kretprobe(event="udp_recvmsg", fn_name="trace_udp_recvmsg_ret")

    csv_file = open(args.output, "w")
    csv_file.write("timestamp,event_type,value_ns,pid\n")

    def handle_event(cpu, data, size):
        event = b["events"].event(data)
        etype = EVENT_TYPES.get(event.event_type, "UNKNOWN")
        csv_file.write(f"{time.time()},{etype},{event.delta_ns},{event.pid}\n")

        # Update live stats
        if event.event_type == 0:
            stats["send_count"] += 1
            stats["send_total_ns"] += event.delta_ns
        elif event.event_type == 1:
            stats["recv_count"] += 1
            stats["recv_total_ns"] += event.delta_ns
        elif event.event_type == 2:
            stats["drop_count"] += 1

    b["events"].open_perf_buffer(handle_event)

    def signal_handler(sig, frame):
        print("\n--- Measurement Summary ---")
        if stats["send_count"] > 0:
            avg_send = stats["send_total_ns"] / stats["send_count"] / 1000
            print(f"  Send events:  {stats['send_count']:>8}  (avg {avg_send:.1f} µs)")
        if stats["recv_count"] > 0:
            avg_recv = stats["recv_total_ns"] / stats["recv_count"] / 1000
            print(f"  Recv events:  {stats['recv_count']:>8}  (avg {avg_recv:.1f} µs)")
        print(f"  Drop events:  {stats['drop_count']:>8}")
        print(f"  Output file:  {args.output}")
        csv_file.close()
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    print(f"Tracing... Output → {args.output}. Ctrl-C to stop.")
    while True:
        b.perf_buffer_poll()


if __name__ == "__main__":
    main()