package engine

import "net"

type L3Info struct {
	Version  uint8
	Protocol uint8
	IHL      uint8
	SrcIP    [4]byte
	DstIP    [4]byte
	Length   uint16
}

func (i L3Info) SrcIPAddr() net.IP { return net.IP(i.SrcIP[:]) }
func (i L3Info) DstIPAddr() net.IP { return net.IP(i.DstIP[:]) }

type TCPInfo struct {
	SrcPort uint16
	DstPort uint16
	Seq     uint32
	Ack     uint32
	HdrLen  uint8
	SYN     bool
	FIN     bool
	RST     bool
	ACK     bool
}

type UDPInfo struct {
	SrcPort uint16
	DstPort uint16
}

func ParseL3(data []byte) (l3 L3Info, transport []byte, ok bool) {
	if len(data) < 20 {
		return
	}
	version := data[0] >> 4
	if version != 4 {
		return
	}
	ihl := data[0] & 0x0F
	if ihl < 5 || len(data) < int(ihl)*4 {
		return
	}
	totalLen := int(uint16(data[2])<<8 | uint16(data[3]))
	if totalLen < int(ihl)*4 || totalLen > len(data) {
		totalLen = len(data)
	}
	return L3Info{
		Version:  4,
		Protocol: data[9],
		IHL:      ihl,
		Length:   uint16(totalLen),
		SrcIP:    [4]byte{data[12], data[13], data[14], data[15]},
		DstIP:    [4]byte{data[16], data[17], data[18], data[19]},
	}, data[ihl*4:totalLen], true
}

func ParseTCP(transport []byte) (TCPInfo, []byte, bool) {
	if len(transport) < 20 {
		return TCPInfo{}, nil, false
	}
	dataOff := uint8(transport[12]>>4) * 4
	if dataOff < 20 || len(transport) < int(dataOff) {
		return TCPInfo{}, nil, false
	}
	flags := transport[13]
	payloadLen := len(transport) - int(dataOff)
	return TCPInfo{
		SrcPort: uint16(transport[0])<<8 | uint16(transport[1]),
		DstPort: uint16(transport[2])<<8 | uint16(transport[3]),
		Seq:     uint32(transport[4])<<24 | uint32(transport[5])<<16 | uint32(transport[6])<<8 | uint32(transport[7]),
		Ack:     uint32(transport[8])<<24 | uint32(transport[9])<<16 | uint32(transport[10])<<8 | uint32(transport[11]),
		HdrLen:  dataOff,
		SYN:     flags&0x02 != 0,
		FIN:     flags&0x01 != 0,
		RST:     flags&0x04 != 0,
		ACK:     flags&0x10 != 0,
	}, transport[dataOff : dataOff+uint8(payloadLen)], true
}

func ParseUDP(transport []byte) (UDPInfo, []byte, bool) {
	if len(transport) < 8 {
		return UDPInfo{}, nil, false
	}
	return UDPInfo{
		SrcPort: uint16(transport[0])<<8 | uint16(transport[1]),
		DstPort: uint16(transport[2])<<8 | uint16(transport[3]),
	}, transport[8:], true
}