package aerotun

import (
	"encoding/binary"

	"github.com/pierrec/lz4/v4"
)

// isEncryptedIPPacket inspects the packet's IP and transport headers to determine
// if the payload is already encrypted or part of a protocol with high entropy
// that cannot be compressed with LZ4.
func isEncryptedIPPacket(b []byte) bool {
	if len(b) < 20 {
		return false
	}
	ver := b[0] >> 4
	if ver == 4 {
		ihl := int(b[0]&0x0f) * 4
		if len(b) < ihl {
			return false
		}
		proto := b[9]
		if proto == 50 { // IPsec ESP
			return true
		}
		if proto == 6 || proto == 17 { // TCP or UDP
			if len(b) < ihl+4 {
				return false
			}
			srcPort := binary.BigEndian.Uint16(b[ihl : ihl+2])
			dstPort := binary.BigEndian.Uint16(b[ihl+2 : ihl+4])
			return isEncryptedPort(srcPort) || isEncryptedPort(dstPort)
		}
	} else if ver == 6 {
		if len(b) < 40 {
			return false
		}
		next := b[6]
		if next == 50 { // IPsec ESP
			return true
		}
		if next == 6 || next == 17 { // TCP or UDP
			if len(b) < 44 {
				return false
			}
			srcPort := binary.BigEndian.Uint16(b[40:42])
			dstPort := binary.BigEndian.Uint16(b[42:44])
			return isEncryptedPort(srcPort) || isEncryptedPort(dstPort)
		}
	}
	return false
}

// isEncryptedPort checks whether the given TCP/UDP port is a well-known
// encrypted or high-entropy transport protocol.
func isEncryptedPort(port uint16) bool {
	switch port {
	case 22, // SSH / SFTP
		443,   // HTTPS, HTTP/3 (QUIC)
		465,   // SMTPS
		500,   // IKE (IPsec)
		853,   // DNS-over-TLS (DoT), DNS-over-QUIC (DoQ)
		993,   // IMAPS
		995,   // POP3S
		1194,  // OpenVPN
		4500,  // IPsec NAT-Traversal
		8443,  // Alternative HTTPS
		51820, // WireGuard
		51830: // AeroTUN encapsulated
		return true
	default:
		return false
	}
}

// adaptiveCompressor manages block compression with fast port bypass and
// adaptive failure backoff, preventing CPU starvation from uncompressible streams.
type adaptiveCompressor struct {
	compBuf   []byte
	skipCount int
	failCount int
}

func newAdaptiveCompressor(mtu int) *adaptiveCompressor {
	if mtu <= 0 {
		mtu = DefaultMTU
	}
	return &adaptiveCompressor{
		compBuf: make([]byte, lz4.CompressBlockBound(mtu)),
	}
}

// Compress attempts to compress the payload using LZ4. It returns the compressed
// slice and true if compression succeeded and saved at least 16 bytes.
// If the packet is unsuitable, encrypted, or during adaptive backoff, it returns
// the original payload and false without wasting CPU cycles.
func (c *adaptiveCompressor) Compress(payload []byte) ([]byte, bool) {
	// Skip small packets (TCP ACKs, keepalives, small DNS queries)
	if len(payload) < 128 {
		return payload, false
	}

	// Fast bypass: skip known encrypted protocols instantly
	if isEncryptedIPPacket(payload) {
		return payload, false
	}

	// Adaptive backoff: if consecutive attempts recently failed, skip without invoking LZ4
	if c.skipCount > 0 {
		c.skipCount--
		return payload, false
	}

	n, err := lz4.CompressBlock(payload, c.compBuf, nil)
	if err == nil && n > 0 && n < len(payload)-16 {
		c.failCount = 0
		return c.compBuf[:n], true
	}

	// Failed to save at least 16 bytes: record failure
	c.failCount++
	if c.failCount >= 16 {
		// Incompressible stream detected: back off for the next 256 packets before probing again
		c.skipCount = 256
		c.failCount = 0
	}

	return payload, false
}