package main

/*
#include <libproc.h>
#include <sys/proc_info.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <stdlib.h>
#include <string.h>

// Define constants that might be macros or hard to use directly
const int C_PROC_PIDLISTFDS = PROC_PIDLISTFDS;
const int C_PROC_PIDFDSOCKETINFO = PROC_PIDFDSOCKETINFO;
const int C_PROX_FDTYPE_SOCKET = PROX_FDTYPE_SOCKET;
const int C_IPPROTO_TCP = IPPROTO_TCP;
const int C_IPPROTO_UDP = IPPROTO_UDP;

// Helper function to access socket info because nested structs can be tricky in Go CGO
uint16_t get_tcp_local_port(struct socket_fdinfo *si) {
    return si->psi.soi_proto.pri_tcp.tcpsi_ini.insi_lport;
}

uint16_t get_udp_local_port(struct socket_fdinfo *si) {
    return si->psi.soi_proto.pri_in.insi_lport;
}

int get_protocol(struct socket_fdinfo *si) {
    return si->psi.soi_protocol;
}
*/
import "C"
import (
	"encoding/binary"
	"fmt"
	"net"
	"os"
	"time"
	"unsafe"
)

const (
	// Protocol constants mirrored from C
	IPPROTO_TCP = C.IPPROTO_TCP
	IPPROTO_UDP = C.IPPROTO_UDP
)

// Scanner provides optimized process scanning capabilities
type Scanner struct {
	pidBuf []C.pid_t
	fdBuf  []byte // Backing store for C.struct_proc_fdinfo array
}

func NewScanner() *Scanner {
	return &Scanner{
		pidBuf: make([]C.pid_t, 4096), // Start with reasonable size
		fdBuf:  make([]byte, 16*1024), // 16KB for FDs
	}
}

// FindProcessConnectingOnPort searches for a process owning the given port (host byte order).
// Returns pid, path, true, total_procs, nil if found.
// Returns error if syscall fails in a way that aborts the scan.
func (s *Scanner) FindProcessConnectingOnPort(protocol int, port uint16) (int, string, bool, int, error) {
	// 1. Get List of PIDs
	// We optimize by trying to reuse s.pidBuf.
	var numPids int
	for {
		// proc_listpids(type, typeinfo, buffer, buffersize)
		// We use type=1 (PROC_ALL_PIDS) or type=2 (PROC_UID_ONLY, if we wanted to filter)
		bufSize := len(s.pidBuf) * int(unsafe.Sizeof(s.pidBuf[0]))
		ret := C.proc_listpids(C.PROC_ALL_PIDS, 0, unsafe.Pointer(&s.pidBuf[0]), C.int(bufSize))

		if ret < 0 {
			return 0, "", false, 0, fmt.Errorf("proc_listpids failed")
		}

		// If buffer was large enough, ret is the number of bytes used.
		bytesReturned := int(ret)
		if bytesReturned < bufSize {
			numPids = bytesReturned / int(unsafe.Sizeof(s.pidBuf[0]))
			break
		}

		// Resize
		newCap := len(s.pidBuf) * 2
		s.pidBuf = make([]C.pid_t, newCap)
	}

	targetPortNet := htons(port)
	cProtocol := C.int(protocol)

	// 2. Iterate PIDs
	for i := 0; i < numPids; i++ {
		pid := s.pidBuf[i]
		if pid <= 0 {
			continue
		}

		matched, err := s.checkPid(pid, cProtocol, targetPortNet)
		if err != nil {
			// ignore specific pid errors, continue scanning
			continue
		}
		if matched {
			// Get Path
			pathBuf := make([]byte, C.PROC_PIDPATHINFO_MAXSIZE)
			retLen := C.proc_pidpath(pid, unsafe.Pointer(&pathBuf[0]), C.uint32_t(len(pathBuf)))
			path := ""
			if retLen > 0 {
				path = C.GoString((*C.char)(unsafe.Pointer(&pathBuf[0])))
			}
			return int(pid), path, true, numPids, nil
		}
	}

	return 0, "", false, numPids, nil
}

func (s *Scanner) checkPid(pid C.pid_t, protocol C.int, targetPortNet uint16) (bool, error) {
	// Optimistic Read of FDs
	fdItemSize := int(C.sizeof_struct_proc_fdinfo)

	// Try with current buffer
	ret := C.proc_pidinfo(pid, C.C_PROC_PIDLISTFDS, 0, unsafe.Pointer(&s.fdBuf[0]), C.int(len(s.fdBuf)))
	if ret <= 0 {
		return false, nil // Process might have exited or access denied
	}

	bytesNeeded := int(ret)
	if bytesNeeded > len(s.fdBuf) {
		// Query actual need
		retSize := C.proc_pidinfo(pid, C.C_PROC_PIDLISTFDS, 0, nil, 0)
		if retSize <= 0 {
			return false, nil
		}

		actualNeeded := int(retSize)
		if actualNeeded > len(s.fdBuf) {
			s.fdBuf = make([]byte, actualNeeded+4096)
		}

		// Retry
		ret = C.proc_pidinfo(pid, C.C_PROC_PIDLISTFDS, 0, unsafe.Pointer(&s.fdBuf[0]), C.int(len(s.fdBuf)))
		if ret <= 0 {
			return false, nil
		}
		bytesNeeded = int(ret)
	}

	numFDs := bytesNeeded / fdItemSize
	basePtr := unsafe.Pointer(&s.fdBuf[0])

	for i := range numFDs {
		itemPtr := (*C.struct_proc_fdinfo)(unsafe.Pointer(uintptr(basePtr) + uintptr(i*fdItemSize)))

		if itemPtr.proc_fdtype != C.C_PROX_FDTYPE_SOCKET {
			continue
		}

		// Get Socket Info
		var sockInfo C.struct_socket_fdinfo
		// proc_pidfdinfo(pid, fd, type, buf, size)
		retSock := C.proc_pidfdinfo(pid, itemPtr.proc_fd, C.C_PROC_PIDFDSOCKETINFO, unsafe.Pointer(&sockInfo), C.int(C.sizeof_struct_socket_fdinfo))
		if int(retSock) != C.sizeof_struct_socket_fdinfo {
			continue
		}

		// Check Protocol
		if C.get_protocol(&sockInfo) != protocol {
			continue
		}

		// Check Port
		var localPort uint16
		if protocol == C.C_IPPROTO_TCP {
			localPort = uint16(C.get_tcp_local_port(&sockInfo))
		} else {
			localPort = uint16(C.get_udp_local_port(&sockInfo))
		}

		if localPort == targetPortNet {
			return true, nil
		}
	}

	return false, nil
}

func htons(v uint16) uint16 {
	// Little Endian (Mac/Intel/ARM) to Big Endian (Network)
	b := [2]byte{}
	binary.BigEndian.PutUint16(b[:], v)
	return *(*uint16)(unsafe.Pointer(&b[0]))
}

// Main function for testing / demonstration
func main() {
	scanner := NewScanner()

	fmt.Println("Starting Go Process Scanner Demo...")

	// ===================================
	// TCP Test (Server)
	// ===================================
	tcpServer, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		fmt.Printf("Failed to listen tcp: %v\n", err)
		os.Exit(1)
	}
	defer tcpServer.Close()

	tcpServerAddr := tcpServer.Addr().(*net.TCPAddr)
	fmt.Printf("TCP Server listening addr %s\n", tcpServerAddr.String())

	go func() {
		conn, _ := tcpServer.Accept()
		defer conn.Close()

		if conn != nil {
			tcpRemoteAddr := conn.RemoteAddr().(*net.TCPAddr)
			fmt.Printf("TCP Client connected from %s\n", tcpRemoteAddr.String())

			checkPort(scanner, IPPROTO_TCP, uint16(tcpRemoteAddr.Port), "TCP Server")
		}
	}()

	// ===================================
	// TCP Test (Client)
	// ===================================
	tcpClient, err := net.Dial("tcp", tcpServer.Addr().String())
	if err != nil {
		fmt.Printf("Failed to connect tcp: %v\n", err)
		os.Exit(1)
	}
	defer tcpClient.Close()

	time.Sleep(time.Microsecond * 100)

	// ===================================
	// UDP Test (Server)
	// ===================================
	udpServer, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		fmt.Printf("Failed to listen udp: %v\n", err)
		os.Exit(1)
	}
	defer udpServer.Close()

	udpServerAddr := udpServer.LocalAddr().(*net.UDPAddr)
	fmt.Printf("UDP Server listening addr %s\n", udpServerAddr.String())

	go func() {
		buf := make([]byte, 1024)
		_, addr, _ := udpServer.ReadFrom(buf)

		if addr != nil {
			udpRemoteAddr := addr.(*net.UDPAddr)
			fmt.Printf("UDP Client connected from %s\n", udpRemoteAddr.String())

			checkPort(scanner, IPPROTO_UDP, uint16(udpRemoteAddr.Port), "UDP Server")
		}
	}()

	// ===================================
	// UDP Test (Client)
	// ===================================
	udpClient, err := net.Dial("udp", udpServer.LocalAddr().String())
	if err != nil {
		fmt.Printf("Failed to dial udp: %v\n", err)
		os.Exit(1)
	}
	defer udpClient.Close()

	udpClient.Write([]byte("ping"))
	time.Sleep(time.Microsecond * 100)
}

func checkPort(scanner *Scanner, protocol int, port uint16, label string) {
	start := time.Now()
	pid, path, found, total, err := scanner.FindProcessConnectingOnPort(protocol, port)
	elapsed := time.Since(start)

	if err != nil {
		fmt.Printf("Error scanning for %s: %v\n", label, err)
		os.Exit(1)
	}

	if found {
		fmt.Printf("[%s] Found PID: %d, Path: %s, Total: %d, Time: %v\n", label, pid, path, total, elapsed)
	} else {
		fmt.Printf("FAILURE: %s process not found. Total: %d\n", label, total)
		os.Exit(1)
	}
}