package engine

import "net"

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

func (i L3Info) SrcIPAddr() net.IP {
	if i.Version == 6 {
		return net.IP(i.SrcIPv6[:])
	}
	return net.IP(i.SrcIP[:])
}

func (i L3Info) DstIPAddr() net.IP {
	if i.Version == 6 {
		return net.IP(i.DstIPv6[:])
	}
	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) < 1 {
		return
	}
	version := data[0] >> 4
	switch version {
	case 4:
		if len(data) < 20 {
			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
	case 6:
		if len(data) < 40 {
			return
		}
		payloadLen := int(uint16(data[4])<<8 | uint16(data[5]))
		totalLen := 40 + payloadLen
		if payloadLen == 0 || totalLen > len(data) {
			totalLen = len(data)
		}
		protocol, tr, ipv6OK := parseIPv6Transport(data, data[6], totalLen)
		if !ipv6OK {
			return
		}
		var srcIP, dstIP [16]byte
		copy(srcIP[:], data[8:24])
		copy(dstIP[:], data[24:40])
		return L3Info{
			Version:  6,
			Protocol: protocol,
			SrcIPv6:  srcIP,
			DstIPv6:  dstIP,
			Length:   uint16(totalLen),
		}, tr, true
	default:
		return
	}
}

func parseIPv6Transport(data []byte, nextHeader uint8, totalLen int) (protocol uint8, transport []byte, ok bool) {
	offset := 40
	proto := nextHeader
	for {
		if offset > totalLen {
			return 0, nil, false
		}
		switch proto {
		case 0, 43, 60: // Hop-by-hop options, Routing, Destination options
			if offset+2 > totalLen {
				return 0, nil, false
			}
			hdrLen := (int(data[offset+1]) + 1) * 8
			if hdrLen < 8 || offset+hdrLen > totalLen {
				return 0, nil, false
			}
			proto = data[offset]
			offset += hdrLen
		case 44: // Fragment
			if offset+8 > totalLen {
				return 0, nil, false
			}
			// Only first fragment carries L4 headers.
			fragOffset := (uint16(data[offset+2])<<8 | uint16(data[offset+3])) >> 3
			if fragOffset != 0 {
				return 0, nil, false
			}
			proto = data[offset]
			offset += 8
		case 51: // Authentication Header
			if offset+2 > totalLen {
				return 0, nil, false
			}
			hdrLen := (int(data[offset+1]) + 2) * 4
			if hdrLen < 8 || offset+hdrLen > totalLen {
				return 0, nil, false
			}
			proto = data[offset]
			offset += hdrLen
		default:
			if offset > totalLen {
				return 0, nil, false
			}
			return proto, data[offset: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 : int(dataOff)+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
}