// Package dhcp implements a minimal DHCP client operating on raw Ethernet frames.
//
// It constructs DHCP DISCOVER/REQUEST as complete Ethernet frames (including
// IP and UDP headers) and parses DHCP OFFER/ACK responses. This allows the
// VPN client to obtain an IP address without an external dhcpcd dependency.
//
// The DHCP exchange is performed over the VPN tunnel: frames are sent via
// the tunnel's WriteFrames and responses are identified from incoming tunnel frames.
package dhcp

import (
	"encoding/binary"
	"fmt"
	"math/rand"
	"net"
	"time"
)

// Route is a classless static route from DHCP option 121/249.
type Route struct {
	Dest    net.IPNet
	Gateway net.IP
}

// Lease holds the result of a successful DHCP exchange.
type Lease struct {
	ClientIP   net.IP
	SubnetMask net.IPMask
	Gateway    net.IP
	DNS        []net.IP
	Routes     []Route // Classless static routes (option 121/249)
	LeaseTime  time.Duration
	ServerIP   net.IP // DHCP server identifier
}

const (
	bootRequest   = 1
	bootReply     = 2
	hwEthernet    = 1
	hwAddrLen     = 6
	dhcpMagic     = 0x63825363
	clientPort    = 68
	serverPort    = 67
	maxDHCPSize   = 576

	// DHCP message types (option 53)
	dhcpDiscover = 1
	dhcpOffer    = 2
	dhcpRequest  = 3
	dhcpAck      = 5
	dhcpNak      = 6

	// DHCP options
	optSubnetMask    = 1
	optRouter        = 3
	optDNS           = 6
	optLeaseTime     = 51
	optMessageType   = 53
	optServerID      = 54
	optRequestedIP        = 50
	optParamRequest       = 55
	optClasslessRoutes    = 121
	optMSClasslessRoutes  = 249
	optEnd                = 255
)

// Client performs DHCP over raw Ethernet frames sent through a VPN tunnel.
type Client struct {
	MAC     net.HardwareAddr // Client MAC address (from TAP device)
	xid     uint32           // Transaction ID
	frameCh chan []byte       // Channel to receive candidate DHCP response frames
}

// NewClient creates a DHCP client for the given MAC address.
func NewClient(mac net.HardwareAddr) *Client {
	return &Client{
		MAC:     mac,
		xid:     rand.Uint32(),
		frameCh: make(chan []byte, 16),
	}
}

// FeedFrame should be called with every Ethernet frame received from the tunnel.
// It checks if the frame is a DHCP response and queues it for processing.
func (c *Client) FeedFrame(frame []byte) {
	if !isDHCPResponse(frame, c.xid) {
		return
	}
	select {
	case c.frameCh <- append([]byte(nil), frame...):
	default:
	}
}

// Run performs a full DHCP exchange (DISCOVER→OFFER→REQUEST→ACK).
// sendFrame is called to transmit each Ethernet frame through the tunnel.
// Returns the lease on success.
func (c *Client) Run(sendFrame func([]byte) error, timeout time.Duration) (*Lease, error) {
	c.xid = rand.Uint32()
	// Send DISCOVER
	disc := c.buildDiscover()
	if err := sendFrame(disc); err != nil {
		return nil, fmt.Errorf("send discover: %w", err)
	}

	// Wait for OFFER
	offer, err := c.waitForType(dhcpOffer, timeout)
	if err != nil {
		return nil, fmt.Errorf("waiting for offer: %w", err)
	}

	offeredIP := net.IP(offer.yiaddr[:])
	serverID := offer.getOptionIP(optServerID)
	if serverID == nil {
		return nil, fmt.Errorf("offer missing server identifier")
	}

	// Send REQUEST
	req := c.buildRequest(offeredIP, serverID)
	if err := sendFrame(req); err != nil {
		return nil, fmt.Errorf("send request: %w", err)
	}

	// Wait for ACK
	ack, err := c.waitForType(dhcpAck, timeout)
	if err != nil {
		return nil, fmt.Errorf("waiting for ack: %w", err)
	}

	lease := &Lease{
		ClientIP:   net.IP(ack.yiaddr[:]).To4(),
		ServerIP:   ack.getOptionIP(optServerID),
		SubnetMask: net.IPMask(ack.getOptionRaw(optSubnetMask)),
		Gateway:    ack.getOptionIP(optRouter),
		DNS:        ack.getOptionIPs(optDNS),
	}
	if lt := ack.getOptionUint32(optLeaseTime); lt > 0 {
		lease.LeaseTime = time.Duration(lt) * time.Second
	}
	if lease.SubnetMask == nil {
		lease.SubnetMask = net.CIDRMask(24, 32)
	}

	// Parse classless static routes (option 121, falls back to 249)
	if routes := parseClasslessRoutes(ack.getOptionRaw(optClasslessRoutes)); len(routes) > 0 {
		lease.Routes = routes
	} else if routes := parseClasslessRoutes(ack.getOptionRaw(optMSClasslessRoutes)); len(routes) > 0 {
		lease.Routes = routes
	}

	return lease, nil
}

// Renew sends a DHCP REQUEST to renew the current lease.
// In RENEWING state: ciaddr = current IP, no requested-IP or server-ID options.
// Returns the renewed lease on ACK, or error on NAK/timeout.
func (c *Client) Renew(currentIP net.IP, sendFrame func([]byte) error, timeout time.Duration) (*Lease, error) {
	c.xid = rand.Uint32()
	var ciaddr [4]byte
	copy(ciaddr[:], currentIP.To4())

	opts := []dhcpOption{
		{optMessageType, []byte{dhcpRequest}},
		{optParamRequest, []byte{optSubnetMask, optRouter, optDNS, optLeaseTime, optClasslessRoutes, optMSClasslessRoutes}},
	}
	frame := c.buildFrame(opts, ciaddr, nil)
	if err := sendFrame(frame); err != nil {
		return nil, fmt.Errorf("send renew: %w", err)
	}

	ack, err := c.waitForType(dhcpAck, timeout)
	if err != nil {
		return nil, fmt.Errorf("waiting for renew ack: %w", err)
	}

	lease := &Lease{
		ClientIP:   net.IP(ack.yiaddr[:]).To4(),
		ServerIP:   ack.getOptionIP(optServerID),
		SubnetMask: net.IPMask(ack.getOptionRaw(optSubnetMask)),
		Gateway:    ack.getOptionIP(optRouter),
		DNS:        ack.getOptionIPs(optDNS),
	}
	if lt := ack.getOptionUint32(optLeaseTime); lt > 0 {
		lease.LeaseTime = time.Duration(lt) * time.Second
	}
	if lease.SubnetMask == nil {
		lease.SubnetMask = net.CIDRMask(24, 32)
	}
	if routes := parseClasslessRoutes(ack.getOptionRaw(optClasslessRoutes)); len(routes) > 0 {
		lease.Routes = routes
	} else if routes := parseClasslessRoutes(ack.getOptionRaw(optMSClasslessRoutes)); len(routes) > 0 {
		lease.Routes = routes
	}

	return lease, nil
}

func (c *Client) waitForType(msgType byte, timeout time.Duration) (*dhcpMsg, error) {
	deadline := time.After(timeout)
	for {
		select {
		case <-deadline:
			return nil, fmt.Errorf("timeout waiting for DHCP message type %d", msgType)
		case frame := <-c.frameCh:
			msg, err := parseDHCPFrame(frame)
			if err != nil {
				continue
			}
			if msg.getOptionByte(optMessageType) == msgType {
				return msg, nil
			}
		}
	}
}

// dhcpMsg is a parsed DHCP message.
type dhcpMsg struct {
	op     byte
	xid    uint32
	yiaddr [4]byte
	siaddr [4]byte
	options []dhcpOption
}

type dhcpOption struct {
	code byte
	data []byte
}

func (m *dhcpMsg) getOption(code byte) *dhcpOption {
	for i := range m.options {
		if m.options[i].code == code {
			return &m.options[i]
		}
	}
	return nil
}

func (m *dhcpMsg) getOptionByte(code byte) byte {
	if o := m.getOption(code); o != nil && len(o.data) >= 1 {
		return o.data[0]
	}
	return 0
}

func (m *dhcpMsg) getOptionIP(code byte) net.IP {
	if o := m.getOption(code); o != nil && len(o.data) >= 4 {
		return net.IP(o.data[:4]).To4()
	}
	return nil
}

func (m *dhcpMsg) getOptionIPs(code byte) []net.IP {
	o := m.getOption(code)
	if o == nil {
		return nil
	}
	var ips []net.IP
	for i := 0; i+3 < len(o.data); i += 4 {
		ips = append(ips, net.IP(o.data[i:i+4]).To4())
	}
	return ips
}

func (m *dhcpMsg) getOptionRaw(code byte) []byte {
	if o := m.getOption(code); o != nil {
		return o.data
	}
	return nil
}

func (m *dhcpMsg) getOptionUint32(code byte) uint32 {
	if o := m.getOption(code); o != nil && len(o.data) >= 4 {
		return binary.BigEndian.Uint32(o.data)
	}
	return 0
}

// buildDiscover creates a full Ethernet frame containing a DHCP DISCOVER.
func (c *Client) buildDiscover() []byte {
	opts := []dhcpOption{
		{optMessageType, []byte{dhcpDiscover}},
		{optParamRequest, []byte{optSubnetMask, optRouter, optDNS, optLeaseTime, optClasslessRoutes, optMSClasslessRoutes}},
	}
	return c.buildFrame(opts, [4]byte{}, nil)
}

// buildRequest creates a full Ethernet frame containing a DHCP REQUEST.
func (c *Client) buildRequest(requestedIP, serverID net.IP) []byte {
	opts := []dhcpOption{
		{optMessageType, []byte{dhcpRequest}},
		{optRequestedIP, requestedIP.To4()},
		{optServerID, serverID.To4()},
		{optParamRequest, []byte{optSubnetMask, optRouter, optDNS, optLeaseTime, optClasslessRoutes, optMSClasslessRoutes}},
	}
	return c.buildFrame(opts, [4]byte{}, serverID.To4())
}

// buildFrame constructs a complete Ethernet/IP/UDP/DHCP frame.
func (c *Client) buildFrame(opts []dhcpOption, ciaddr [4]byte, _ net.IP) []byte {
	// Build DHCP payload
	dhcp := make([]byte, 240) // fixed DHCP header
	dhcp[0] = bootRequest
	dhcp[1] = hwEthernet
	dhcp[2] = hwAddrLen
	binary.BigEndian.PutUint32(dhcp[4:8], c.xid)
	copy(dhcp[12:16], ciaddr[:])
	copy(dhcp[28:34], c.MAC)
	binary.BigEndian.PutUint32(dhcp[236:240], dhcpMagic)

	// Append options
	for _, o := range opts {
		dhcp = append(dhcp, o.code, byte(len(o.data)))
		dhcp = append(dhcp, o.data...)
	}
	dhcp = append(dhcp, optEnd)

	// Pad to minimum size
	for len(dhcp) < 300 {
		dhcp = append(dhcp, 0)
	}

	// Build UDP header (8 bytes)
	udpLen := uint16(8 + len(dhcp))
	udp := make([]byte, 8)
	binary.BigEndian.PutUint16(udp[0:2], clientPort)
	binary.BigEndian.PutUint16(udp[2:4], serverPort)
	binary.BigEndian.PutUint16(udp[4:6], udpLen)
	// checksum = 0 (optional for IPv4 UDP)

	// Build IPv4 header (20 bytes, no options)
	ipLen := uint16(20 + len(udp) + len(dhcp))
	ip := make([]byte, 20)
	ip[0] = 0x45 // version=4, IHL=5
	binary.BigEndian.PutUint16(ip[2:4], ipLen)
	binary.BigEndian.PutUint16(ip[6:8], 0) // flags=0, fragment=0
	ip[8] = 64                              // TTL
	ip[9] = 17                              // protocol = UDP
	copy(ip[12:16], net.IPv4zero.To4())     // src = 0.0.0.0
	copy(ip[16:20], net.IPv4bcast.To4())    // dst = 255.255.255.255
	// IP header checksum
	binary.BigEndian.PutUint16(ip[10:12], ipChecksum(ip))

	// Build Ethernet header (14 bytes)
	eth := make([]byte, 14)
	copy(eth[0:6], net.HardwareAddr{0xff, 0xff, 0xff, 0xff, 0xff, 0xff}) // dst = broadcast
	copy(eth[6:12], c.MAC)                                                // src
	binary.BigEndian.PutUint16(eth[12:14], 0x0800)                        // EtherType = IPv4

	frame := make([]byte, 0, len(eth)+len(ip)+len(udp)+len(dhcp))
	frame = append(frame, eth...)
	frame = append(frame, ip...)
	frame = append(frame, udp...)
	frame = append(frame, dhcp...)
	return frame
}

// isDHCPResponse checks if an Ethernet frame is a DHCP response (UDP from port 67 to 68)
// matching our transaction ID.
func isDHCPResponse(frame []byte, xid uint32) bool {
	if len(frame) < 14+20+8+240 { // eth + ip + udp + dhcp minimum
		return false
	}
	// EtherType must be IPv4
	if binary.BigEndian.Uint16(frame[12:14]) != 0x0800 {
		return false
	}
	ipStart := 14
	// IP protocol must be UDP (17)
	if frame[ipStart+9] != 17 {
		return false
	}
	ihl := int(frame[ipStart]&0x0f) * 4
	udpStart := ipStart + ihl
	if len(frame) < udpStart+8 {
		return false
	}
	srcPort := binary.BigEndian.Uint16(frame[udpStart : udpStart+2])
	dstPort := binary.BigEndian.Uint16(frame[udpStart+2 : udpStart+4])
	if srcPort != serverPort || dstPort != clientPort {
		return false
	}
	dhcpStart := udpStart + 8
	if len(frame) < dhcpStart+240 {
		return false
	}
	// Check xid
	pktXid := binary.BigEndian.Uint32(frame[dhcpStart+4 : dhcpStart+8])
	return pktXid == xid
}

// parseDHCPFrame extracts a DHCP message from a full Ethernet frame.
func parseDHCPFrame(frame []byte) (*dhcpMsg, error) {
	ipStart := 14
	ihl := int(frame[ipStart]&0x0f) * 4
	dhcpStart := ipStart + ihl + 8 // skip IP + UDP headers
	data := frame[dhcpStart:]

	if len(data) < 240 {
		return nil, fmt.Errorf("dhcp packet too short")
	}

	msg := &dhcpMsg{
		op:  data[0],
		xid: binary.BigEndian.Uint32(data[4:8]),
	}
	copy(msg.yiaddr[:], data[16:20])
	copy(msg.siaddr[:], data[20:24])

	// Parse options (after magic cookie at offset 236)
	magic := binary.BigEndian.Uint32(data[236:240])
	if magic != dhcpMagic {
		return nil, fmt.Errorf("bad DHCP magic: %x", magic)
	}

	opts := data[240:]
	for len(opts) > 0 {
		code := opts[0]
		if code == optEnd {
			break
		}
		if code == 0 { // padding
			opts = opts[1:]
			continue
		}
		if len(opts) < 2 {
			break
		}
		length := int(opts[1])
		if len(opts) < 2+length {
			break
		}
		msg.options = append(msg.options, dhcpOption{
			code: code,
			data: append([]byte(nil), opts[2:2+length]...),
		})
		opts = opts[2+length:]
	}
	return msg, nil
}

// parseClasslessRoutes decodes RFC 3442 classless static routes.
// Encoding: [mask_width] [significant dest octets...] [gateway 4 bytes]
// e.g. /0 = 1+4 bytes (default route), /25 = 1+4+4 bytes
func parseClasslessRoutes(data []byte) []Route {
	if len(data) == 0 {
		return nil
	}
	var routes []Route
	for len(data) > 0 {
		if len(data) < 1 {
			break
		}
		maskWidth := int(data[0])
		data = data[1:]
		if maskWidth > 32 {
			break
		}
		// Number of significant destination octets
		nOctets := (maskWidth + 7) / 8
		if len(data) < nOctets+4 {
			break
		}
		var dest [4]byte
		copy(dest[:], data[:nOctets])
		data = data[nOctets:]
		gw := net.IP(append([]byte(nil), data[:4]...)).To4()
		data = data[4:]

		routes = append(routes, Route{
			Dest: net.IPNet{
				IP:   net.IP(dest[:]).To4(),
				Mask: net.CIDRMask(maskWidth, 32),
			},
			Gateway: gw,
		})
	}
	return routes
}

func ipChecksum(header []byte) uint16 {
	var sum uint32
	for i := 0; i < len(header)-1; i += 2 {
		sum += uint32(binary.BigEndian.Uint16(header[i : i+2]))
	}
	if len(header)%2 == 1 {
		sum += uint32(header[len(header)-1]) << 8
	}
	for sum > 0xffff {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return ^uint16(sum)
}