package main

import (
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"net"
	"os"
	"regexp"
	"sort"
	"strings"
	"syscall"
	"time"
)

const (
	ethPAll      = 0x0003
	ethPIP       = 0x0800
	ethP8021Q    = 0x8100
	ethP8021AD   = 0x88A8
	packetOut    = 4
	protoICMP    = 1
	protoTCP     = 6
	protoUDP     = 17
	maxFrameSize = 65535
)

var vmInterfaceRE = regexp.MustCompile(`(?:tap|fwbr|fwln|fwpr)(\d+)`)
var vmInterfaceDetailRE = regexp.MustCompile(`(?:tap|fwbr|fwln|fwpr)(\d+)i(\d+)`)

type flowKey struct {
	VMID            string
	NIC             string
	Interface       string
	Direction       string
	SourceIP        string
	DestinationIP   string
	Protocol        string
	SourcePort      int
	DestinationPort int
}

type flowValue struct {
	SourceIP        string `json:"source_ip"`
	DestinationIP   string `json:"destination_ip"`
	Protocol        string `json:"protocol"`
	SourcePort      *int   `json:"source_port,omitempty"`
	DestinationPort *int   `json:"destination_port,omitempty"`
	Packets         uint64 `json:"packets"`
	Bytes           uint64 `json:"bytes"`
	VMID            string `json:"vmid,omitempty"`
	NIC             string `json:"nic,omitempty"`
	Interface       string `json:"interface,omitempty"`
	Direction       string `json:"direction,omitempty"`
	State           string `json:"state"`
	Collector       string `json:"collector"`
}

type diagnostics struct {
	AttachMode         string   `json:"attach_mode"`
	InterfacesRequested []string `json:"interfaces_requested"`
	InterfacesAttached  []string `json:"interfaces_attached"`
	Errors             []string `json:"errors"`
}

type payload struct {
	Flows       []flowValue `json:"flows"`
	Diagnostics diagnostics `json:"diagnostics"`
}

func htons(value uint16) int {
	return int((value<<8)&0xff00 | value>>8)
}

func intPtr(value int) *int {
	if value == 0 {
		return nil
	}
	return &value
}

func protocolName(value byte) string {
	switch value {
	case protoICMP:
		return "icmp"
	case protoTCP:
		return "tcp"
	case protoUDP:
		return "udp"
	default:
		return ""
	}
}

func parsePacket(packet []byte) (flowKey, int, bool) {
	var key flowKey
	if len(packet) < 34 {
		return key, 0, false
	}
	offset := 12
	ethType := binary.BigEndian.Uint16(packet[offset : offset+2])
	offset = 14
	for ethType == ethP8021Q || ethType == ethP8021AD {
		if len(packet) < offset+4 {
			return key, 0, false
		}
		ethType = binary.BigEndian.Uint16(packet[offset+2 : offset+4])
		offset += 4
	}
	if ethType != ethPIP || len(packet) < offset+20 {
		return key, 0, false
	}
	versionIHL := packet[offset]
	version := versionIHL >> 4
	ihl := int(versionIHL&0x0f) * 4
	if version != 4 || ihl < 20 || len(packet) < offset+ihl {
		return key, 0, false
	}
	totalLength := int(binary.BigEndian.Uint16(packet[offset+2 : offset+4]))
	protocol := protocolName(packet[offset+9])
	if protocol == "" {
		return key, 0, false
	}
	key.SourceIP = net.IP(packet[offset+12 : offset+16]).String()
	key.DestinationIP = net.IP(packet[offset+16 : offset+20]).String()
	key.Protocol = protocol
	transportOffset := offset + ihl
	if protocol == "tcp" || protocol == "udp" {
		if len(packet) < transportOffset+4 {
			return key, 0, false
		}
		key.SourcePort = int(binary.BigEndian.Uint16(packet[transportOffset : transportOffset+2]))
		key.DestinationPort = int(binary.BigEndian.Uint16(packet[transportOffset+2 : transportOffset+4]))
	} else if protocol == "icmp" && len(packet) >= transportOffset+2 {
		key.SourcePort = int(packet[transportOffset])
		key.DestinationPort = int(packet[transportOffset+1])
	}
	if totalLength <= 0 {
		totalLength = len(packet) - offset
	}
	return key, totalLength, true
}

func interfaceMeta(name string) (string, string) {
	vmid := ""
	nic := "0"
	if match := vmInterfaceRE.FindStringSubmatch(name); len(match) > 1 {
		vmid = match[1]
	}
	if match := vmInterfaceDetailRE.FindStringSubmatch(name); len(match) > 2 {
		nic = match[2]
	}
	return vmid, nic
}

func setFd(fd int, set *syscall.FdSet) {
	set.Bits[fd/64] |= 1 << uint(fd%64)
}

func isSet(fd int, set *syscall.FdSet) bool {
	return set.Bits[fd/64]&(1<<uint(fd%64)) != 0
}

func openSocket(interfaceName string) (int, error) {
	iface, err := net.InterfaceByName(interfaceName)
	if err != nil {
		return -1, err
	}
	fd, err := syscall.Socket(syscall.AF_PACKET, syscall.SOCK_RAW, htons(ethPAll))
	if err != nil {
		return -1, err
	}
	addr := &syscall.SockaddrLinklayer{Protocol: htons(ethPAll), Ifindex: iface.Index}
	if err := syscall.Bind(fd, addr); err != nil {
		_ = syscall.Close(fd)
		return -1, err
	}
	if err := syscall.SetNonblock(fd, true); err != nil {
		_ = syscall.Close(fd)
		return -1, err
	}
	return fd, nil
}

func collect(interfaceNames []string, duration time.Duration, limit int) payload {
	result := payload{
		Diagnostics: diagnostics{
			AttachMode:          "af_packet_raw_socket",
			InterfacesRequested: interfaceNames,
			Errors:              []string{},
		},
	}
	type socketInfo struct {
		name string
		fd   int
	}
	sockets := []socketInfo{}
	for _, name := range interfaceNames {
		if strings.TrimSpace(name) == "" {
			continue
		}
		fd, err := openSocket(strings.TrimSpace(name))
		if err != nil {
			result.Diagnostics.Errors = append(result.Diagnostics.Errors, fmt.Sprintf("%s: %v", name, err))
			continue
		}
		sockets = append(sockets, socketInfo{name: strings.TrimSpace(name), fd: fd})
		result.Diagnostics.InterfacesAttached = append(result.Diagnostics.InterfacesAttached, strings.TrimSpace(name))
	}
	defer func() {
		for _, socket := range sockets {
			_ = syscall.Close(socket.fd)
		}
	}()
	if len(sockets) == 0 {
		return result
	}

	fdToSocket := map[int]socketInfo{}
	maxFd := 0
	for _, socket := range sockets {
		fdToSocket[socket.fd] = socket
		if socket.fd > maxFd {
			maxFd = socket.fd
		}
	}
	flows := map[flowKey]*flowValue{}
	deadline := time.Now().Add(duration)
	buffer := make([]byte, maxFrameSize)
	for time.Now().Before(deadline) {
		var readfds syscall.FdSet
		for _, socket := range sockets {
			setFd(socket.fd, &readfds)
		}
		timeout := syscall.NsecToTimeval(int64(250 * time.Millisecond))
		_, err := syscall.Select(maxFd+1, &readfds, nil, nil, &timeout)
		if err != nil && err != syscall.EINTR {
			result.Diagnostics.Errors = append(result.Diagnostics.Errors, err.Error())
			break
		}
		for fd, socket := range fdToSocket {
			if !isSet(fd, &readfds) {
				continue
			}
			n, from, err := syscall.Recvfrom(fd, buffer, 0)
			if err != nil {
				if err != syscall.EAGAIN && err != syscall.EWOULDBLOCK {
					result.Diagnostics.Errors = append(result.Diagnostics.Errors, fmt.Sprintf("%s: %v", socket.name, err))
				}
				continue
			}
			key, bytes, ok := parsePacket(buffer[:n])
			if !ok {
				continue
			}
			key.Interface = socket.name
			key.VMID, key.NIC = interfaceMeta(socket.name)
			key.Direction = "ingress"
			if link, ok := from.(*syscall.SockaddrLinklayer); ok && link.Pkttype == packetOut {
				key.Direction = "egress"
			}
			current := flows[key]
			if current == nil {
				current = &flowValue{
					SourceIP:        key.SourceIP,
					DestinationIP:   key.DestinationIP,
					Protocol:        key.Protocol,
					SourcePort:      intPtr(key.SourcePort),
					DestinationPort: intPtr(key.DestinationPort),
					VMID:            key.VMID,
					NIC:             key.NIC,
					Interface:       key.Interface,
					Direction:       key.Direction,
					State:           "observed",
					Collector:       "ebpf-helper",
				}
				flows[key] = current
			}
			current.Packets++
			current.Bytes += uint64(bytes)
			if len(flows) >= limit {
				break
			}
		}
		if len(flows) >= limit {
			break
		}
	}
	for _, flow := range flows {
		result.Flows = append(result.Flows, *flow)
	}
	sort.Slice(result.Flows, func(i, j int) bool {
		if result.Flows[i].Bytes == result.Flows[j].Bytes {
			return result.Flows[i].Packets > result.Flows[j].Packets
		}
		return result.Flows[i].Bytes > result.Flows[j].Bytes
	})
	if len(result.Flows) > limit {
		result.Flows = result.Flows[:limit]
	}
	return result
}

func main() {
	jsonOutput := flag.Bool("json", false, "print JSON output")
	limit := flag.Int("limit", 1500, "maximum unique flows")
	durationSeconds := flag.Int("duration", 10, "collection duration in seconds")
	interfaces := flag.String("interfaces", "", "comma-separated interface names")
	flag.Parse()
	if !*jsonOutput {
		fmt.Fprintln(os.Stderr, "only --json output is supported")
		os.Exit(2)
	}
	if *limit <= 0 {
		*limit = 1500
	}
	if *durationSeconds <= 0 {
		*durationSeconds = 10
	}
	names := []string{}
	for _, name := range strings.Split(*interfaces, ",") {
		if strings.TrimSpace(name) != "" {
			names = append(names, strings.TrimSpace(name))
		}
	}
	result := collect(names, time.Duration(*durationSeconds)*time.Second, *limit)
	data, err := json.Marshal(result)
	if err != nil {
		fmt.Fprintf(os.Stderr, "json marshal failed: %v\n", err)
		os.Exit(1)
	}
	fmt.Println(string(data))
}