package divert

import (
	"encoding/binary"
	"errors"
	"net"
)

// IPv4TCPInfo is what we extract from a raw IPv4+TCP packet for our
// per-flow mapping table.
type IPv4TCPInfo struct {
	SrcIP, DstIP     net.IP
	SrcPort, DstPort uint16
}

// ParseIPv4TCP reads the IPv4 + TCP header pair out of an outbound
// captured packet and returns the addressing info. Does NOT mutate
// the buffer.
//
// Errors when:
//   - buffer too short to contain a full IPv4+TCP header (40 bytes)
//   - IP version is not 4
//   - IP protocol is not 6 (TCP)
func ParseIPv4TCP(b []byte) (*IPv4TCPInfo, error) {
	if len(b) < 40 {
		return nil, errors.New("packet shorter than IPv4+TCP minimum")
	}
	if b[0]>>4 != 4 {
		return nil, errors.New("not IPv4")
	}
	ihl := int(b[0]&0x0f) * 4
	if ihl < 20 || len(b) < ihl+20 {
		return nil, errors.New("IPv4 IHL invalid or buffer truncated")
	}
	if b[9] != 6 {
		return nil, errors.New("not TCP")
	}
	src := net.IPv4(b[12], b[13], b[14], b[15])
	dst := net.IPv4(b[16], b[17], b[18], b[19])
	srcPort := binary.BigEndian.Uint16(b[ihl : ihl+2])
	dstPort := binary.BigEndian.Uint16(b[ihl+2 : ihl+4])
	return &IPv4TCPInfo{
		SrcIP:   src,
		DstIP:   dst,
		SrcPort: srcPort,
		DstPort: dstPort,
	}, nil
}

// RewriteDst mutates b in-place to set dst IP and port, then
// recomputes both the IP header checksum and the TCP checksum.
//
// Returns the same errors as ParseIPv4TCP for malformed input.
func RewriteDst(b []byte, ip net.IP, port uint16) error {
	if _, err := ParseIPv4TCP(b); err != nil {
		return err
	}
	v4 := ip.To4()
	if v4 == nil {
		return errors.New("dst must be IPv4")
	}
	ihl := int(b[0]&0x0f) * 4

	// Set dst IP
	copy(b[16:20], v4)
	// Set dst port
	binary.BigEndian.PutUint16(b[ihl+2:ihl+4], port)

	// Recompute IP checksum (clear → compute → write big-endian)
	b[10], b[11] = 0, 0
	cs := ipChecksum(b[:ihl])
	b[10] = byte(cs >> 8)
	b[11] = byte(cs & 0xff)

	// Recompute TCP checksum (clear → compute → write)
	b[ihl+16], b[ihl+17] = 0, 0
	cs = tcpChecksum(b[:ihl], b[ihl:])
	b[ihl+16] = byte(cs >> 8)
	b[ihl+17] = byte(cs & 0xff)

	return nil
}

// SwapAndSetDstPort applies the canonical streamdump-style NAT-redirect
// rewrite: swap IPv4 src/dst, set TCP dst port to newDstPort. Keeps
// the original TCP src port (so the listener sees a unique RemoteAddr
// it can use to look up the flow). Recomputes both checksums.
//
// Use this on the FORWARD path (outbound from target process →
// remote). After this rewrite, set addr.Outbound=0 and reinject —
// the packet looks like remote → local on the inbound path, lands at
// the listener.
func SwapAndSetDstPort(b []byte, newDstPort uint16) error {
	if _, err := ParseIPv4TCP(b); err != nil {
		return err
	}
	ihl := int(b[0]&0x0f) * 4

	// Swap src ↔ dst IPv4 (bytes 12..15 ↔ 16..19)
	var src, dst [4]byte
	copy(src[:], b[12:16])
	copy(dst[:], b[16:20])
	copy(b[12:16], dst[:])
	copy(b[16:20], src[:])

	// Set TCP dst port; src port unchanged.
	binary.BigEndian.PutUint16(b[ihl+2:ihl+4], newDstPort)

	// Recompute IP checksum
	b[10], b[11] = 0, 0
	cs := ipChecksum(b[:ihl])
	b[10] = byte(cs >> 8)
	b[11] = byte(cs & 0xff)

	// Recompute TCP checksum
	b[ihl+16], b[ihl+17] = 0, 0
	cs = tcpChecksum(b[:ihl], b[ihl:])
	b[ihl+16] = byte(cs >> 8)
	b[ihl+17] = byte(cs & 0xff)
	return nil
}

// SwapAndSetSrcPort applies the canonical streamdump-style return-path
// rewrite: swap IPv4 src/dst, set TCP src port to newSrcPort (the
// original target port the client expects to see, e.g. 443). Keeps
// the original TCP dst port (which is the client's ephemeral port).
//
// Use this on the RETURN path (listener → client). After this rewrite,
// set addr.Outbound=0 and reinject — the packet looks like remote →
// local on the inbound path, matches the client's connect() pair, and
// the client socket accepts the response as if from the real target.
func SwapAndSetSrcPort(b []byte, newSrcPort uint16) error {
	if _, err := ParseIPv4TCP(b); err != nil {
		return err
	}
	ihl := int(b[0]&0x0f) * 4

	// Swap src ↔ dst IPv4
	var src, dst [4]byte
	copy(src[:], b[12:16])
	copy(dst[:], b[16:20])
	copy(b[12:16], dst[:])
	copy(b[16:20], src[:])

	// Set TCP src port; dst port unchanged.
	binary.BigEndian.PutUint16(b[ihl:ihl+2], newSrcPort)

	// Recompute IP checksum
	b[10], b[11] = 0, 0
	cs := ipChecksum(b[:ihl])
	b[10] = byte(cs >> 8)
	b[11] = byte(cs & 0xff)

	// Recompute TCP checksum
	b[ihl+16], b[ihl+17] = 0, 0
	cs = tcpChecksum(b[:ihl], b[ihl:])
	b[ihl+16] = byte(cs >> 8)
	b[ihl+17] = byte(cs & 0xff)
	return nil
}

// ipChecksum is the standard 16-bit one's-complement sum over the IP
// header (RFC 791). The "checksum field" must be zeroed before calling.
func ipChecksum(hdr []byte) uint16 {
	var sum uint32
	for i := 0; i+1 < len(hdr); i += 2 {
		sum += uint32(hdr[i])<<8 | uint32(hdr[i+1])
	}
	if len(hdr)%2 == 1 {
		sum += uint32(hdr[len(hdr)-1]) << 8
	}
	for sum>>16 != 0 {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return ^uint16(sum)
}

// tcpChecksum implements the RFC 793 pseudo-header checksum.
// ipHdr must include src+dst addresses; tcpSeg is the full TCP header
// + payload. The "checksum field" inside tcpSeg must be zeroed.
func tcpChecksum(ipHdr, tcpSeg []byte) uint16 {
	var sum uint32
	// Pseudo-header: src(4) dst(4) zero(1) proto(1) tcp_len(2)
	for i := 12; i <= 18; i += 2 {
		sum += uint32(ipHdr[i])<<8 | uint32(ipHdr[i+1])
	}
	sum += uint32(6) // TCP protocol
	tcpLen := uint32(len(tcpSeg))
	sum += tcpLen
	// TCP segment
	for i := 0; i+1 < len(tcpSeg); i += 2 {
		sum += uint32(tcpSeg[i])<<8 | uint32(tcpSeg[i+1])
	}
	if len(tcpSeg)%2 == 1 {
		sum += uint32(tcpSeg[len(tcpSeg)-1]) << 8
	}
	for sum>>16 != 0 {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return ^uint16(sum)
}

// IPv4UDPInfo is what we extract from a raw IPv4+UDP packet for our
// per-flow mapping table.
type IPv4UDPInfo struct {
	SrcIP, DstIP     net.IP
	SrcPort, DstPort uint16
	IHL              int // IPv4 header length in bytes
	UDPLen           uint16
}

// ParseIPv4UDP reads the IPv4 + UDP header pair out of an outbound
// captured packet and returns the addressing info. Does NOT mutate
// the buffer.
//
// Errors when:
//   - buffer too short to contain a full IPv4+UDP header (28 bytes)
//   - IP version is not 4
//   - IP protocol is not 17 (UDP)
func ParseIPv4UDP(b []byte) (*IPv4UDPInfo, error) {
	if len(b) < 28 {
		return nil, errors.New("packet shorter than IPv4+UDP minimum")
	}
	if b[0]>>4 != 4 {
		return nil, errors.New("not IPv4")
	}
	ihl := int(b[0]&0x0f) * 4
	if ihl < 20 || len(b) < ihl+8 {
		return nil, errors.New("IPv4 IHL invalid or buffer truncated")
	}
	if b[9] != 17 {
		return nil, errors.New("not UDP")
	}
	src := net.IPv4(b[12], b[13], b[14], b[15])
	dst := net.IPv4(b[16], b[17], b[18], b[19])
	srcPort := binary.BigEndian.Uint16(b[ihl : ihl+2])
	dstPort := binary.BigEndian.Uint16(b[ihl+2 : ihl+4])
	udpLen := binary.BigEndian.Uint16(b[ihl+4 : ihl+6])
	return &IPv4UDPInfo{
		SrcIP:   src,
		DstIP:   dst,
		SrcPort: srcPort,
		DstPort: dstPort,
		IHL:     ihl,
		UDPLen:  udpLen,
	}, nil
}

// SwapUDPAndSetDstPort applies the canonical streamdump-style swap to
// a UDP packet: swap IPv4 src/dst, set UDP dst port to newDstPort.
// Keeps the original UDP src port. Recomputes IP and UDP checksums.
//
// (For UDP, swap+reinject is generally NOT used by drover — the
// engine's diverterLoop "consumes" target UDP packets and forwards
// them through the SOCKS5 UDP relay directly. This helper is here for
// completeness/symmetry with the TCP swap helpers and for tests.)
func SwapUDPAndSetDstPort(b []byte, newDstPort uint16) error {
	if _, err := ParseIPv4UDP(b); err != nil {
		return err
	}
	ihl := int(b[0]&0x0f) * 4

	// Swap src ↔ dst IPv4
	var src, dst [4]byte
	copy(src[:], b[12:16])
	copy(dst[:], b[16:20])
	copy(b[12:16], dst[:])
	copy(b[16:20], src[:])

	// Set UDP dst port
	binary.BigEndian.PutUint16(b[ihl+2:ihl+4], newDstPort)

	// Recompute IP checksum
	b[10], b[11] = 0, 0
	cs := ipChecksum(b[:ihl])
	b[10] = byte(cs >> 8)
	b[11] = byte(cs & 0xff)

	// Recompute UDP checksum (offset ihl+6,ihl+7 inside UDP header)
	udpLen := int(binary.BigEndian.Uint16(b[ihl+4 : ihl+6]))
	if ihl+udpLen > len(b) {
		udpLen = len(b) - ihl
	}
	b[ihl+6], b[ihl+7] = 0, 0
	cs = udpChecksum(b[:ihl], b[ihl:ihl+udpLen])
	// Zero is "no checksum" in IPv4 UDP. RFC 768 says when the
	// computed checksum is zero, transmit it as 0xFFFF instead.
	if cs == 0 {
		cs = 0xFFFF
	}
	b[ihl+6] = byte(cs >> 8)
	b[ihl+7] = byte(cs & 0xff)
	return nil
}

// SwapUDPAndSetSrcPort applies the canonical streamdump-style return-
// path swap to a UDP packet: swap IPv4 src/dst, set UDP src port to
// newSrcPort (the original target port the client expects to see).
// Recomputes IP and UDP checksums. (Symmetric counterpart to the TCP
// helper; not currently used by the engine for the same reason as
// SwapUDPAndSetDstPort, but exists for tests/parity.)
func SwapUDPAndSetSrcPort(b []byte, newSrcPort uint16) error {
	if _, err := ParseIPv4UDP(b); err != nil {
		return err
	}
	ihl := int(b[0]&0x0f) * 4

	// Swap src ↔ dst IPv4
	var src, dst [4]byte
	copy(src[:], b[12:16])
	copy(dst[:], b[16:20])
	copy(b[12:16], dst[:])
	copy(b[16:20], src[:])

	// Set UDP src port
	binary.BigEndian.PutUint16(b[ihl:ihl+2], newSrcPort)

	// Recompute IP checksum
	b[10], b[11] = 0, 0
	cs := ipChecksum(b[:ihl])
	b[10] = byte(cs >> 8)
	b[11] = byte(cs & 0xff)

	// Recompute UDP checksum
	udpLen := int(binary.BigEndian.Uint16(b[ihl+4 : ihl+6]))
	if ihl+udpLen > len(b) {
		udpLen = len(b) - ihl
	}
	b[ihl+6], b[ihl+7] = 0, 0
	cs = udpChecksum(b[:ihl], b[ihl:ihl+udpLen])
	if cs == 0 {
		cs = 0xFFFF
	}
	b[ihl+6] = byte(cs >> 8)
	b[ihl+7] = byte(cs & 0xff)
	return nil
}

// BuildIPv4UDPInbound fabricates an IPv4+UDP packet for reinjection
// as inbound (return path from upstream relay → Discord). Used by the
// UDPProxy after the SOCKS5 relay sends back a response: we construct
// a synthetic packet that looks like remote_endpoint → local_IP and
// reinject it via WinDivert with addr.Outbound=0.
//
//	src      → original Discord destination (the UDP server)
//	dst      → local LAN IP we bound on
//	srcPort  → original destination port (e.g. 50007)
//	dstPort  → Discord's ephemeral src port (so the kernel matches the
//	           connect()-bound socket)
//
// The returned slice owns its own backing storage; callers may pass
// it directly to (*Handle).Send.
func BuildIPv4UDPInbound(srcIP, dstIP net.IP, srcPort, dstPort uint16, payload []byte) ([]byte, error) {
	src := srcIP.To4()
	dst := dstIP.To4()
	if src == nil || dst == nil {
		return nil, errors.New("BuildIPv4UDPInbound: src/dst must be IPv4")
	}
	if len(payload)+28 > 0xFFFF {
		return nil, errors.New("BuildIPv4UDPInbound: payload too large for IPv4 datagram")
	}

	totalLen := 20 + 8 + len(payload)
	buf := make([]byte, totalLen)

	// IPv4 header (20 bytes, IHL=5, no options)
	buf[0] = 0x45 // version=4, IHL=5
	buf[1] = 0x00 // DSCP/ECN
	binary.BigEndian.PutUint16(buf[2:4], uint16(totalLen))
	binary.BigEndian.PutUint16(buf[4:6], 0) // ID
	binary.BigEndian.PutUint16(buf[6:8], 0) // flags + frag
	buf[8] = 64                              // TTL
	buf[9] = 17                              // protocol = UDP
	// checksum at [10..11] left zero for now
	copy(buf[12:16], src)
	copy(buf[16:20], dst)

	// UDP header (8 bytes)
	binary.BigEndian.PutUint16(buf[20:22], srcPort)
	binary.BigEndian.PutUint16(buf[22:24], dstPort)
	binary.BigEndian.PutUint16(buf[24:26], uint16(8+len(payload))) // UDP length
	// UDP checksum at [26..27] left zero for now

	// Payload
	copy(buf[28:], payload)

	// Recompute IP checksum
	cs := ipChecksum(buf[:20])
	buf[10] = byte(cs >> 8)
	buf[11] = byte(cs & 0xff)

	// Recompute UDP checksum (over pseudo-header + UDP segment)
	cs = udpChecksum(buf[:20], buf[20:])
	if cs == 0 {
		cs = 0xFFFF // RFC 768: 0 means "checksum disabled", send 0xFFFF instead
	}
	buf[26] = byte(cs >> 8)
	buf[27] = byte(cs & 0xff)

	return buf, nil
}

// udpChecksum implements the RFC 768 pseudo-header checksum for IPv4
// UDP. ipHdr must include src+dst addresses; udpSeg is the full UDP
// header + payload (UDP "length" field already set; checksum field
// inside udpSeg must be zeroed).
//
// IPv4 UDP checksum is technically OPTIONAL — a sender may transmit
// 0 to indicate "no checksum". We always compute one since most
// modern stacks (and Discord) expect a valid checksum.
func udpChecksum(ipHdr, udpSeg []byte) uint16 {
	var sum uint32
	// Pseudo-header: src(4) dst(4) zero(1) proto(1) udp_len(2)
	for i := 12; i <= 18; i += 2 {
		sum += uint32(ipHdr[i])<<8 | uint32(ipHdr[i+1])
	}
	sum += uint32(17) // UDP protocol
	udpLen := uint32(len(udpSeg))
	sum += udpLen
	// UDP segment (header + payload)
	for i := 0; i+1 < len(udpSeg); i += 2 {
		sum += uint32(udpSeg[i])<<8 | uint32(udpSeg[i+1])
	}
	if len(udpSeg)%2 == 1 {
		sum += uint32(udpSeg[len(udpSeg)-1]) << 8
	}
	for sum>>16 != 0 {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return ^uint16(sum)
}