// xdp_lb.c
#include <stdio.h>
#include <stdlib.h>
#include <bpf/libbpf.h>
#include <bpf/bpf.h>
#include <net/if.h>
#include <linux/if_link.h>
#include "xdp_lb.skel.h"  // This header is auto-generated by bpftool

#define IFACE "eth0"  // Replace with your network interface

static int set_up_backends(struct xdp_lb_bpf *skel) {
    __u32 backend1 = htonl(0xC0A80102);  // 192.168.1.2
    __u32 backend2 = htonl(0xC0A80103);  // 192.168.1.3
    __u32 key = 0;

    if (bpf_map_update_elem(bpf_map__fd(skel->maps.backends), &key, &backend1, BPF_ANY) < 0) {
        fprintf(stderr, "Failed to update backend 1\n");
        return -1;
    }

    key = 1;
    if (bpf_map_update_elem(bpf_map__fd(skel->maps.backends), &key, &backend2, BPF_ANY) < 0) {
        fprintf(stderr, "Failed to update backend 2\n");
        return -1;
    }

    return 0;
}

int main() {
    struct xdp_lb_bpf *skel;
    int err, ifindex;

    // Load and verify the eBPF skeleton
    skel = xdp_lb_bpf__open();
    if (!skel) {
        fprintf(stderr, "Failed to open and load skeleton\n");
        return 1;
    }

    // Load eBPF program
    err = xdp_lb_bpf__load(skel);
    if (err) {
        fprintf(stderr, "Failed to load BPF program: %d\n", err);
        return 1;
    }

    // Set up the backend IP addresses
    if (set_up_backends(skel) < 0) {
        fprintf(stderr, "Failed to set up backend IP addresses\n");
        return 1;
    }

    // Get interface index
    ifindex = if_nametoindex(IFACE);
    if (ifindex == 0) {
        perror("if_nametoindex");
        return 1;
    }

    // Attach the XDP program
    err = bpf_xdp_attach(ifindex, bpf_program__fd(skel->progs.xdp_load_balancer), XDP_FLAGS_SKB_MODE, NULL);
    if (err) {
        fprintf(stderr, "Failed to attach XDP program: %d\n", err);
        return 1;
    }

    printf("XDP Load Balancer is running on interface %s...\n", IFACE);
    sleep(60);  // Keep running for 60 seconds

    // Detach the XDP program before exiting
    err = bpf_xdp_detach(ifindex, XDP_FLAGS_SKB_MODE, NULL);
    if (err) {
        fprintf(stderr, "Failed to detach XDP program: %d\n", err);
        return 1;
    }

    // Clean up
    xdp_lb_bpf__destroy(skel);

    return 0;
}