#!/usr/bin/env python3
"""Emit bridge ingress raw packets plus egress/drop telemetry via tc/eBPF."""

from __future__ import annotations

import argparse
import base64
import ctypes as ct
import hashlib
import ipaddress
import json
import signal
import socket
import subprocess
import sys
from pathlib import Path
from typing import Iterable

from pyroute2 import IPRoute

try:
    from bcc import BPF  # type: ignore
except Exception as exc:  # pragma: no cover - depends on host runtime
    dist_packages = [
        Path("/usr/lib/python3/dist-packages"),
        Path("/usr/lib64/python3/dist-packages"),
    ]
    for candidate in dist_packages:
        candidate_str = str(candidate)
        if candidate.is_dir() and candidate_str not in sys.path:
            sys.path.append(candidate_str)
    try:
        from bcc import BPF  # type: ignore
    except Exception:
        print(f"Failed to import python3-bpfcc: {exc}", file=sys.stderr, flush=True)
        raise

from src.utilities.packet_mark import packet_id_from_mark, verdict_from_mark


EVENT_INGRESS = 1
EVENT_EGRESS = 2
EVENT_DROP = 3
INGRESS_PARENT = "ffff:fff2"
EGRESS_PARENT = "ffff:fff3"
INGRESS_FILTER_HANDLE = ":20"
EGRESS_FILTER_HANDLE = ":30"

IDENTITY_FIELDS = (
    "src_mac",
    "dst_mac",
    "eth_type_raw",
    "vlan_id",
    "src_ip",
    "dst_ip",
    "protocol_raw",
    "src_port",
    "dst_port",
    "length",
)

BPF_SOURCE = r"""
#include <uapi/linux/ptrace.h>
#include <uapi/linux/pkt_cls.h>
#include <linux/bpf.h>
#include <linux/skbuff.h>
#include <linux/netdevice.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/ipv6.h>
#include <linux/in.h>
#include <linux/tcp.h>
#include <linux/udp.h>
#include <linux/if_arp.h>

#define EVENT_INGRESS 1
#define EVENT_EGRESS  2
#define EVENT_DROP    3

struct vlan_hdr_t {
  __be16 h_vlan_TCI;
  __be16 h_vlan_encapsulated_proto;
};

struct arp_eth_ipv4_t {
  __u8 sha[6];
  __u8 spa[4];
  __u8 tha[6];
  __u8 tpa[4];
};

struct event_t {
  __u64 ts_ns;
  __u32 skb_mark;
  __u32 length;
  __u32 reason;
  __u32 ifindex;
  __u16 eth_type_raw;
  __u16 vlan_id;
  __u16 src_port;
  __u16 dst_port;
  __u32 protocol_raw;
  __u8 event_type;
  __u8 ip_version;
  __u8 reserved[2];
  unsigned char src_mac[6];
  unsigned char dst_mac[6];
  unsigned char src_ip[16];
  unsigned char dst_ip[16];
};

BPF_PERF_OUTPUT(ingress_events);
BPF_PERF_OUTPUT(meta_events);

static __always_inline __u32 ensure_packet_mark(struct __sk_buff *skb) {
  __u32 next = skb->mark;

  if (next) {
    return next;
  }

  next = (__u32)(bpf_ktime_get_ns() & 0x0FFFFFFF);
  if (!next) {
    next = 1;
  }

  skb->mark = next;
  return next;
}

static __always_inline int parse_l3_l4_direct(struct event_t *event, void *data, void *data_end) {
  struct ethhdr *eth = data;
  __be16 eth_proto;
  void *l3;

  if ((void *)(eth + 1) > data_end) {
    return 0;
  }

  __builtin_memcpy(event->src_mac, eth->h_source, ETH_ALEN);
  __builtin_memcpy(event->dst_mac, eth->h_dest, ETH_ALEN);
  eth_proto = eth->h_proto;
  l3 = eth + 1;

  if (eth_proto == htons(ETH_P_8021Q) || eth_proto == htons(ETH_P_8021AD)) {
    struct vlan_hdr_t *vlan = l3;
    if ((void *)(vlan + 1) > data_end) {
      return 0;
    }
    event->vlan_id = ntohs(vlan->h_vlan_TCI) & 0x0fff;
    eth_proto = vlan->h_vlan_encapsulated_proto;
    l3 = vlan + 1;
  }

  event->eth_type_raw = ntohs(eth_proto);

  if (eth_proto == htons(ETH_P_ARP)) {
    struct arphdr *arph = l3;
    struct arp_eth_ipv4_t *body = (void *)(arph + 1);
    if ((void *)(body + 1) > data_end) {
      return 1;
    }
    if (arph->ar_hrd == htons(ARPHRD_ETHER) && arph->ar_pro == htons(ETH_P_IP) &&
        arph->ar_hln == ETH_ALEN && arph->ar_pln == 4) {
      __builtin_memcpy(event->src_ip, body->spa, 4);
      __builtin_memcpy(event->dst_ip, body->tpa, 4);
      event->ip_version = 4;
    }
    return 1;
  }

  if (eth_proto == htons(ETH_P_IP)) {
    struct iphdr *iph = l3;
    if ((void *)(iph + 1) > data_end) {
      return 0;
    }

    event->ip_version = 4;
    event->protocol_raw = iph->protocol;
    __builtin_memcpy(event->src_ip, &iph->saddr, 4);
    __builtin_memcpy(event->dst_ip, &iph->daddr, 4);

    if (iph->protocol == IPPROTO_TCP) {
      struct tcphdr *tcph = (void *)iph + (iph->ihl * 4);
      if ((void *)(tcph + 1) > data_end) {
        return 1;
      }
      event->src_port = ntohs(tcph->source);
      event->dst_port = ntohs(tcph->dest);
    } else if (iph->protocol == IPPROTO_UDP) {
      struct udphdr *udph = (void *)iph + (iph->ihl * 4);
      if ((void *)(udph + 1) > data_end) {
        return 1;
      }
      event->src_port = ntohs(udph->source);
      event->dst_port = ntohs(udph->dest);
    }
    return 1;
  }

  if (eth_proto == htons(ETH_P_IPV6)) {
    struct ipv6hdr *ip6h = l3;
    if ((void *)(ip6h + 1) > data_end) {
      return 0;
    }

    event->ip_version = 6;
    event->protocol_raw = ip6h->nexthdr;
    __builtin_memcpy(event->src_ip, &ip6h->saddr, 16);
    __builtin_memcpy(event->dst_ip, &ip6h->daddr, 16);

    if (ip6h->nexthdr == IPPROTO_TCP) {
      struct tcphdr *tcph = (void *)(ip6h + 1);
      if ((void *)(tcph + 1) > data_end) {
        return 1;
      }
      event->src_port = ntohs(tcph->source);
      event->dst_port = ntohs(tcph->dest);
    } else if (ip6h->nexthdr == IPPROTO_UDP) {
      struct udphdr *udph = (void *)(ip6h + 1);
      if ((void *)(udph + 1) > data_end) {
        return 1;
      }
      event->src_port = ntohs(udph->source);
      event->dst_port = ntohs(udph->dest);
    }
    return 1;
  }

  return 1;
}

static __always_inline int parse_skb_linear(struct event_t *event, struct sk_buff *skb) {
  unsigned char *head = NULL;
  __u16 mac_header = 0;
  __u16 network_header = 0;
  __u16 transport_header = 0;

  if (!skb) {
    return 0;
  }

  bpf_probe_read_kernel(&head, sizeof(head), &skb->head);
  bpf_probe_read_kernel(&mac_header, sizeof(mac_header), &skb->mac_header);
  bpf_probe_read_kernel(&network_header, sizeof(network_header), &skb->network_header);
  bpf_probe_read_kernel(&transport_header, sizeof(transport_header), &skb->transport_header);
  bpf_probe_read_kernel(&event->length, sizeof(event->length), &skb->len);
  bpf_probe_read_kernel(&event->skb_mark, sizeof(event->skb_mark), &skb->mark);

  if (!head) {
    return 0;
  }

  struct ethhdr eth = {};
  unsigned char *eth_ptr = head + mac_header;
  bpf_probe_read_kernel(&eth, sizeof(eth), eth_ptr);
  __builtin_memcpy(event->src_mac, eth.h_source, ETH_ALEN);
  __builtin_memcpy(event->dst_mac, eth.h_dest, ETH_ALEN);

  __be16 eth_proto = eth.h_proto;
  unsigned char *l3_ptr = head + network_header;
  if (eth_proto == htons(ETH_P_8021Q) || eth_proto == htons(ETH_P_8021AD)) {
    struct vlan_hdr_t vlan = {};
    bpf_probe_read_kernel(&vlan, sizeof(vlan), eth_ptr + sizeof(struct ethhdr));
    event->vlan_id = ntohs(vlan.h_vlan_TCI) & 0x0fff;
    eth_proto = vlan.h_vlan_encapsulated_proto;
  }
  event->eth_type_raw = ntohs(eth_proto);

  if (eth_proto == htons(ETH_P_ARP)) {
    struct arphdr arph = {};
    struct arp_eth_ipv4_t arp_body = {};
    bpf_probe_read_kernel(&arph, sizeof(arph), l3_ptr);
    if (arph.ar_hrd == htons(ARPHRD_ETHER) && arph.ar_pro == htons(ETH_P_IP) &&
        arph.ar_hln == ETH_ALEN && arph.ar_pln == 4) {
      bpf_probe_read_kernel(&arp_body, sizeof(arp_body), l3_ptr + sizeof(struct arphdr));
      __builtin_memcpy(event->src_ip, arp_body.spa, 4);
      __builtin_memcpy(event->dst_ip, arp_body.tpa, 4);
      event->ip_version = 4;
    }
    return 1;
  }

  if (eth_proto == htons(ETH_P_IP)) {
    struct iphdr iph = {};
    bpf_probe_read_kernel(&iph, sizeof(iph), l3_ptr);
    event->ip_version = 4;
    event->protocol_raw = iph.protocol;
    bpf_probe_read_kernel(event->src_ip, 4, &iph.saddr);
    bpf_probe_read_kernel(event->dst_ip, 4, &iph.daddr);

    if (iph.protocol == IPPROTO_TCP) {
      struct tcphdr tcph = {};
      bpf_probe_read_kernel(&tcph, sizeof(tcph), head + transport_header);
      event->src_port = ntohs(tcph.source);
      event->dst_port = ntohs(tcph.dest);
    } else if (iph.protocol == IPPROTO_UDP) {
      struct udphdr udph = {};
      bpf_probe_read_kernel(&udph, sizeof(udph), head + transport_header);
      event->src_port = ntohs(udph.source);
      event->dst_port = ntohs(udph.dest);
    }
    return 1;
  }

  if (eth_proto == htons(ETH_P_IPV6)) {
    struct ipv6hdr ip6h = {};
    bpf_probe_read_kernel(&ip6h, sizeof(ip6h), l3_ptr);
    event->ip_version = 6;
    event->protocol_raw = ip6h.nexthdr;
    __builtin_memcpy(event->src_ip, &ip6h.saddr, 16);
    __builtin_memcpy(event->dst_ip, &ip6h.daddr, 16);

    if (ip6h.nexthdr == IPPROTO_TCP) {
      struct tcphdr tcph = {};
      bpf_probe_read_kernel(&tcph, sizeof(tcph), head + transport_header);
      event->src_port = ntohs(tcph.source);
      event->dst_port = ntohs(tcph.dest);
    } else if (ip6h.nexthdr == IPPROTO_UDP) {
      struct udphdr udph = {};
      bpf_probe_read_kernel(&udph, sizeof(udph), head + transport_header);
      event->src_port = ntohs(udph.source);
      event->dst_port = ntohs(udph.dest);
    }
    return 1;
  }

  return 1;
}

int handle_ingress(struct __sk_buff *skb) {
  struct event_t event = {};
  void *data = (void *)(long)skb->data;
  void *data_end = (void *)(long)skb->data_end;

  event.ts_ns = bpf_ktime_get_ns();
  event.event_type = EVENT_INGRESS;
  event.ifindex = skb->ifindex;
  event.length = skb->len;
  event.skb_mark = ensure_packet_mark(skb);

  if (!event.skb_mark) {
    return TC_ACT_OK;
  }
  if (!parse_l3_l4_direct(&event, data, data_end)) {
    return TC_ACT_OK;
  }

  ingress_events.perf_submit_skb(skb, skb->len, &event, sizeof(event));
  return TC_ACT_OK;
}

int handle_egress(struct __sk_buff *skb) {
  struct event_t event = {};
  void *data = (void *)(long)skb->data;
  void *data_end = (void *)(long)skb->data_end;

  event.ts_ns = bpf_ktime_get_ns();
  event.event_type = EVENT_EGRESS;
  event.ifindex = skb->ifindex;
  event.length = skb->len;
  event.skb_mark = skb->mark;

  if (!event.skb_mark) {
    return TC_ACT_OK;
  }
  if (!parse_l3_l4_direct(&event, data, data_end)) {
    return TC_ACT_OK;
  }

  meta_events.perf_submit(skb, &event, sizeof(event));
  return TC_ACT_OK;
}

TRACEPOINT_PROBE(skb, kfree_skb) {
  struct sk_buff *skb = (struct sk_buff *)args->skbaddr;
  struct event_t event = {};
  struct net_device *dev = NULL;

  event.ts_ns = bpf_ktime_get_ns();
  event.event_type = EVENT_DROP;
  event.reason = args->reason;

  if (!skb) {
    return 0;
  }

  bpf_probe_read_kernel(&dev, sizeof(dev), &skb->dev);
  if (!dev) {
    return 0;
  }
  bpf_probe_read_kernel(&event.ifindex, sizeof(event.ifindex), &dev->ifindex);
  if (!parse_skb_linear(&event, skb)) {
    return 0;
  }
  if (!event.skb_mark) {
    return 0;
  }

  meta_events.perf_submit(args, &event, sizeof(event));
  return 0;
}
"""


class Event(ct.Structure):
    _fields_ = [
        ("ts_ns", ct.c_ulonglong),
        ("skb_mark", ct.c_uint),
        ("length", ct.c_uint),
        ("reason", ct.c_uint),
        ("ifindex", ct.c_uint),
        ("eth_type_raw", ct.c_ushort),
        ("vlan_id", ct.c_ushort),
        ("src_port", ct.c_ushort),
        ("dst_port", ct.c_ushort),
        ("protocol_raw", ct.c_uint),
        ("event_type", ct.c_ubyte),
        ("ip_version", ct.c_ubyte),
        ("reserved", ct.c_ubyte * 2),
        ("src_mac", ct.c_ubyte * 6),
        ("dst_mac", ct.c_ubyte * 6),
        ("src_ip", ct.c_ubyte * 16),
        ("dst_ip", ct.c_ubyte * 16),
    ]


TARGET_INTERFACES: set[str] = set()
IPR: IPRoute | None = None


def _run_checked(cmd: list[str]) -> None:
    subprocess.run(cmd, check=True, capture_output=True, text=True)


def _ifname_from_index(ifindex: int) -> str | None:
    if ifindex <= 0:
        return None
    try:
        return socket.if_indextoname(ifindex)
    except OSError:
        return None


def _mac_to_str(value: Iterable[int]) -> str:
    return ":".join(f"{byte:02x}" for byte in value)


def _ip_to_str(ip_version: int, raw: Iterable[int]) -> str | None:
    data = bytes(raw)
    if ip_version == 4:
        try:
            return str(ipaddress.IPv4Address(data[:4]))
        except ipaddress.AddressValueError:
            return None
    if ip_version == 6:
        try:
            return str(ipaddress.IPv6Address(data[:16]))
        except ipaddress.AddressValueError:
            return None
    return None


def _build_packet_uid(payload: dict[str, object]) -> str:
    normalized = []
    for field in IDENTITY_FIELDS:
        value = payload.get(field)
        normalized.append("" if value is None else str(value))
    return hashlib.sha1("|".join(normalized).encode("utf-8")).hexdigest()


def _event_name(value: int) -> str:
    return {EVENT_INGRESS: "ingress", EVENT_EGRESS: "egress", EVENT_DROP: "drop"}.get(value, "unknown")


def _reason_name(reason: int) -> str:
    return f"skb_drop_reason_{reason}"


def _build_payload(event: Event) -> dict[str, object] | None:
    iface = _ifname_from_index(int(event.ifindex))
    if iface not in TARGET_INTERFACES:
        return None

    payload: dict[str, object] = {
        "event_type": _event_name(int(event.event_type)),
        "iface": iface,
        "skb_mark": int(event.skb_mark) or None,
        "length": int(event.length),
        "src_mac": _mac_to_str(event.src_mac),
        "dst_mac": _mac_to_str(event.dst_mac),
        "eth_type_raw": int(event.eth_type_raw) or None,
        "vlan_id": int(event.vlan_id) or None,
        "src_ip": _ip_to_str(int(event.ip_version), event.src_ip),
        "dst_ip": _ip_to_str(int(event.ip_version), event.dst_ip),
        "protocol_raw": int(event.protocol_raw) or None,
        "src_port": int(event.src_port) or None,
        "dst_port": int(event.dst_port) or None,
        "reason": _reason_name(int(event.reason)) if int(event.event_type) == EVENT_DROP else None,
        "reason_code": int(event.reason) if int(event.event_type) == EVENT_DROP else None,
    }

    packet_id = packet_id_from_mark(payload.get("skb_mark"))
    if packet_id:
        payload["packet_id"] = packet_id
        payload["correlation_key"] = f"pid:{packet_id}"
        payload["correlation_source"] = "kernel_mark"
        verdict_hint = verdict_from_mark(payload.get("skb_mark"))
        if verdict_hint:
            payload["verdict_hint"] = verdict_hint
    else:
        payload["packet_uid"] = _build_packet_uid(payload)
        payload["correlation_key"] = f"uid:{payload['packet_uid']}"
        payload["correlation_source"] = "legacy_hash"
    return payload


def _emit_ingress_event(cpu: int, data: int, size: int) -> None:
    del cpu
    event = ct.cast(data, ct.POINTER(Event)).contents
    payload = _build_payload(event)
    if payload is None:
        return

    raw_size = size - ct.sizeof(Event)
    if raw_size > 0:
        raw = ct.string_at(data + ct.sizeof(Event), min(raw_size, int(event.length)))
        payload["raw_b64"] = base64.b64encode(raw).decode("ascii")

    print(json.dumps(payload, separators=(",", ":")), flush=True)


def _emit_meta_event(cpu: int, data: int, size: int) -> None:
    del cpu, size
    event = ct.cast(data, ct.POINTER(Event)).contents
    payload = _build_payload(event)
    if payload is None:
        return
    print(json.dumps(payload, separators=(",", ":")), flush=True)


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="tc/eBPF bridge telemetry collector")
    parser.add_argument("--ifaces", required=True, help="Comma-separated list of interfaces to instrument")
    parser.add_argument("--build-dir", required=True, help="Directory for compiled tc BPF objects")
    return parser.parse_args()


def _sigterm(_signum: int, _frame: object) -> None:
    raise KeyboardInterrupt


def _json_safe(value: object) -> object:
    if isinstance(value, bytes):
        return value.decode("utf-8", "replace")
    return value


def _ensure_clean_clsact(iface: str) -> None:
    subprocess.run(["tc", "qdisc", "del", "dev", iface, "clsact"], check=False, capture_output=True, text=True)
    _run_checked(["tc", "qdisc", "add", "dev", iface, "clsact"])


def _attach_tc_programs(bpf: BPF, ifaces: list[str], build_dir: str) -> tuple[str, str]:
    global IPR
    del build_dir
    ingress_fn = bpf.load_func("handle_ingress", BPF.SCHED_CLS)
    egress_fn = bpf.load_func("handle_egress", BPF.SCHED_CLS)

    ipr = IPRoute()
    try:
        for iface in ifaces:
            matches = ipr.link_lookup(ifname=iface)
            if not matches:
                raise RuntimeError(f"Interface not found: {iface}")
            ifindex = matches[0]
            _ensure_clean_clsact(iface)
            ipr.tc(
                "add-filter",
                "bpf",
                ifindex,
                INGRESS_FILTER_HANDLE,
                fd=ingress_fn.fd,
                name=ingress_fn.name,
                parent=INGRESS_PARENT,
                classid=1,
                direct_action=True,
            )
            ipr.tc(
                "add-filter",
                "bpf",
                ifindex,
                EGRESS_FILTER_HANDLE,
                fd=egress_fn.fd,
                name=egress_fn.name,
                parent=EGRESS_PARENT,
                classid=1,
                direct_action=True,
            )
    except Exception:
        for iface in ifaces:
            subprocess.run(["tc", "qdisc", "del", "dev", iface, "clsact"], check=False, capture_output=True, text=True)
        if ipr is not None:
            ipr.close()
            IPR = None
        raise

    return ingress_fn.name, egress_fn.name


def _cleanup_tc(ifaces: Iterable[str]) -> None:
    for iface in ifaces:
        subprocess.run(["tc", "qdisc", "del", "dev", iface, "clsact"], check=False, capture_output=True, text=True)
    global IPR
    if IPR is not None:
        try:
            IPR.close()
        finally:
            IPR = None


def main() -> int:
    args = _parse_args()
    global TARGET_INTERFACES
    TARGET_INTERFACES = {iface.strip() for iface in args.ifaces.split(",") if iface.strip()}
    if not TARGET_INTERFACES:
        print("No interfaces provided", file=sys.stderr, flush=True)
        return 1

    signal.signal(signal.SIGTERM, _sigterm)
    signal.signal(signal.SIGINT, _sigterm)

    bpf = BPF(text=BPF_SOURCE)
    ingress_prog_name = ""
    egress_prog_name = ""
    try:
        ingress_prog_name, egress_prog_name = _attach_tc_programs(bpf, sorted(TARGET_INTERFACES), args.build_dir)
        print(
            json.dumps(
                {
                    "status": "collector_started",
                    "ifaces": sorted(TARGET_INTERFACES),
                    "ingress_program": _json_safe(ingress_prog_name),
                    "egress_program": _json_safe(egress_prog_name),
                    "build_dir": str(args.build_dir),
                },
                separators=(",", ":"),
            ),
            flush=True,
        )

        bpf["ingress_events"].open_perf_buffer(_emit_ingress_event, page_cnt=256)
        bpf["meta_events"].open_perf_buffer(_emit_meta_event, page_cnt=128)
        while True:
            bpf.perf_buffer_poll()
    except KeyboardInterrupt:
        return 0
    finally:
        _cleanup_tc(sorted(TARGET_INTERFACES))


if __name__ == "__main__":
    sys.exit(main())