// Package dhcp4 reads and writes DHCPv4 messages (RFC 2131) and their
// options (RFC 2132), and holds the catalog of options Steward knows by
// name. It does no I/O.
package dhcp4

import (
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"net/netip"
	"strings"
)

// Ports.
const (
	ServerPort = 67
	ClientPort = 68
)

// Op codes.
const (
	BootRequest = 1
	BootReply   = 2
)

// MessageType is option 53.
type MessageType byte

// Message types (RFC 2132 section 9.6).
const (
	Discover MessageType = 1
	Offer    MessageType = 2
	Request  MessageType = 3
	Decline  MessageType = 4
	Ack      MessageType = 5
	Nak      MessageType = 6
	Release  MessageType = 7
	Inform   MessageType = 8
)

func (t MessageType) String() string {
	switch t {
	case Discover:
		return "DHCPDISCOVER"
	case Offer:
		return "DHCPOFFER"
	case Request:
		return "DHCPREQUEST"
	case Decline:
		return "DHCPDECLINE"
	case Ack:
		return "DHCPACK"
	case Nak:
		return "DHCPNAK"
	case Release:
		return "DHCPRELEASE"
	case Inform:
		return "DHCPINFORM"
	}
	return fmt.Sprintf("DHCP type %d", byte(t))
}

// Option codes the server itself uses.
const (
	OptPad           = 0
	OptSubnetMask    = 1
	OptRouters       = 3
	OptDNSServers    = 6
	OptHostName      = 12
	OptDomainName    = 15
	OptBroadcast     = 28
	OptRequestedIP   = 50
	OptLeaseTime     = 51
	OptOverload      = 52
	OptMessageType   = 53
	OptServerID      = 54
	OptParamRequest  = 55
	OptMessage       = 56
	OptMaxMessage    = 57
	OptRenewalTime   = 58
	OptRebindingTime = 59
	OptVendorClass   = 60
	OptClientID      = 61
	OptClientFQDN    = 81
	OptEnd           = 255
)

// FlagBroadcast is the broadcast bit of the flags field: the client can't
// receive unicast before it has an address.
const FlagBroadcast = 0x8000

var magicCookie = []byte{99, 130, 83, 99}

// headerLen is the fixed part of a message, up to the magic cookie.
const headerLen = 236

// MinLen is the shortest message to send: BOOTP's 300 bytes, which some
// old clients and relays insist on.
const MinLen = 300

// Option is one option: its code and its data.
type Option struct {
	Code byte
	Data []byte
}

// Packet is a DHCPv4 message.
type Packet struct {
	Op     byte
	HType  byte
	HLen   byte
	Hops   byte
	XID    uint32
	Secs   uint16
	Flags  uint16
	CIAddr netip.Addr
	YIAddr netip.Addr
	SIAddr netip.Addr
	GIAddr netip.Addr
	CHAddr [16]byte
	SName  string
	File   string
	// Options in the order they were read or are to be written. Options
	// that appear more than once are joined (RFC 3396).
	Options []Option
}

var zero4 = netip.AddrFrom4([4]byte{})

// Parse reads a message.
func Parse(b []byte) (*Packet, error) {
	if len(b) < headerLen+4 {
		return nil, fmt.Errorf("message too short (%d bytes)", len(b))
	}
	p := &Packet{
		Op: b[0], HType: b[1], HLen: b[2], Hops: b[3],
		XID:    binary.BigEndian.Uint32(b[4:8]),
		Secs:   binary.BigEndian.Uint16(b[8:10]),
		Flags:  binary.BigEndian.Uint16(b[10:12]),
		CIAddr: addr4(b[12:16]), YIAddr: addr4(b[16:20]), SIAddr: addr4(b[20:24]), GIAddr: addr4(b[24:28]),
	}
	copy(p.CHAddr[:], b[28:44])
	if p.HLen > 16 {
		return nil, fmt.Errorf("hardware address length %d is over 16", p.HLen)
	}
	if string(b[236:240]) != string(magicCookie) {
		return nil, errors.New("no DHCP magic cookie (a BOOTP message?)")
	}
	opts, err := parseOptions(b[240:])
	if err != nil {
		return nil, err
	}
	// Option overload (52): more options in the file and sname fields.
	sname, file := b[44:108], b[108:236]
	if ov := find(opts, OptOverload); len(ov) == 1 {
		if ov[0]&1 != 0 {
			more, err := parseOptions(file)
			if err != nil {
				return nil, fmt.Errorf("options in the file field: %w", err)
			}
			opts = append(opts, more...)
			file = nil
		}
		if ov[0]&2 != 0 {
			more, err := parseOptions(sname)
			if err != nil {
				return nil, fmt.Errorf("options in the sname field: %w", err)
			}
			opts = append(opts, more...)
			sname = nil
		}
	}
	p.SName, p.File = cString(sname), cString(file)
	p.Options = join(opts)
	return p, nil
}

func addr4(b []byte) netip.Addr { return netip.AddrFrom4([4]byte(b[:4])) }

func cString(b []byte) string {
	if i := strings.IndexByte(string(b), 0); i >= 0 {
		b = b[:i]
	}
	return string(b)
}

func parseOptions(b []byte) ([]Option, error) {
	var out []Option
	for i := 0; i < len(b); {
		code := b[i]
		i++
		switch code {
		case OptPad:
			continue
		case OptEnd:
			return out, nil
		}
		if i >= len(b) {
			return nil, fmt.Errorf("option %d has no length", code)
		}
		n := int(b[i])
		i++
		if i+n > len(b) {
			return nil, fmt.Errorf("option %d is cut short (%d of %d bytes)", code, len(b)-i, n)
		}
		out = append(out, Option{Code: code, Data: append([]byte(nil), b[i:i+n]...)})
		i += n
	}
	// Some clients leave out the end option.
	return out, nil
}

// join concatenates repeated options into one (RFC 3396), keeping the
// position of the first.
func join(opts []Option) []Option {
	var out []Option
	seen := map[byte]int{}
	for _, o := range opts {
		if i, ok := seen[o.Code]; ok {
			out[i].Data = append(out[i].Data, o.Data...)
			continue
		}
		seen[o.Code] = len(out)
		out = append(out, o)
	}
	return out
}

func find(opts []Option, code byte) []byte {
	for _, o := range opts {
		if o.Code == code {
			return o.Data
		}
	}
	return nil
}

// Marshal writes the message. Options longer than 255 bytes are split
// (RFC 3396); the message is padded to MinLen.
func (p *Packet) Marshal() []byte {
	b := make([]byte, headerLen, MinLen+64)
	b[0], b[1], b[2], b[3] = p.Op, p.HType, p.HLen, p.Hops
	binary.BigEndian.PutUint32(b[4:8], p.XID)
	binary.BigEndian.PutUint16(b[8:10], p.Secs)
	binary.BigEndian.PutUint16(b[10:12], p.Flags)
	put4(b[12:16], p.CIAddr)
	put4(b[16:20], p.YIAddr)
	put4(b[20:24], p.SIAddr)
	put4(b[24:28], p.GIAddr)
	copy(b[28:44], p.CHAddr[:])
	copy(b[44:107], p.SName)
	copy(b[108:235], p.File)
	b = append(b, magicCookie...)
	for _, o := range p.Options {
		data := o.Data
		for first := true; first || len(data) > 0; first = false {
			n := min(len(data), 255)
			b = append(b, o.Code, byte(n))
			b = append(b, data[:n]...)
			data = data[n:]
		}
	}
	b = append(b, OptEnd)
	for len(b) < MinLen {
		b = append(b, OptPad)
	}
	return b
}

// Len is the length Marshal would produce.
func (p *Packet) Len() int {
	n := headerLen + 4 + 1
	for _, o := range p.Options {
		n += len(o.Data) + 2*max(1, (len(o.Data)+254)/255)
	}
	return max(n, MinLen)
}

func put4(b []byte, a netip.Addr) {
	if a.Is4() {
		v := a.As4()
		copy(b, v[:])
	}
}

// Get returns an option's data (nil when absent).
func (p *Packet) Get(code byte) []byte { return find(p.Options, code) }

// Has reports whether the option is present.
func (p *Packet) Has(code byte) bool {
	for _, o := range p.Options {
		if o.Code == code {
			return true
		}
	}
	return false
}

// Set adds or replaces an option.
func (p *Packet) Set(code byte, data []byte) {
	for i, o := range p.Options {
		if o.Code == code {
			p.Options[i].Data = data
			return
		}
	}
	p.Options = append(p.Options, Option{Code: code, Data: data})
}

// Type is the message type (0 when the option is missing or malformed).
func (p *Packet) Type() MessageType {
	if d := p.Get(OptMessageType); len(d) == 1 {
		return MessageType(d[0])
	}
	return 0
}

// IPOption returns an option holding one IPv4 address.
func (p *Packet) IPOption(code byte) (netip.Addr, bool) {
	d := p.Get(code)
	if len(d) != 4 {
		return netip.Addr{}, false
	}
	return addr4(d), true
}

// MAC is the client's hardware address, for Ethernet-like types.
func (p *Packet) MAC() net.HardwareAddr {
	n := int(p.HLen)
	if n == 0 || n > 16 {
		return nil
	}
	return net.HardwareAddr(append([]byte(nil), p.CHAddr[:n]...))
}

// ClientID is option 61 in lower-case hex ("" when absent).
func (p *Packet) ClientID() string {
	if d := p.Get(OptClientID); len(d) > 0 {
		return fmt.Sprintf("%x", d)
	}
	return ""
}

// HostName is the name the client gives: option 12, or else the name in
// the client FQDN option (81), without its domain.
func (p *Packet) HostName() string {
	if d := p.Get(OptHostName); len(d) > 0 {
		return strings.TrimRight(string(d), "\x00")
	}
	if d := p.Get(OptClientFQDN); len(d) > 3 {
		name := d[3:]
		if d[0]&0x04 != 0 { // E bit: canonical wire format
			if labels, err := decodeNames(name); err == nil && len(labels) > 0 {
				return strings.SplitN(labels[0], ".", 2)[0]
			}
			return ""
		}
		return strings.SplitN(strings.TrimRight(string(name), "\x00"), ".", 2)[0]
	}
	return ""
}

// ParamRequest is the client's parameter request list (option 55).
func (p *Packet) ParamRequest() []byte { return p.Get(OptParamRequest) }

// MaxMessageSize is the largest reply the client accepts (option 57),
// or 576 (RFC 2131's minimum) when it doesn't say.
func (p *Packet) MaxMessageSize() int {
	if d := p.Get(OptMaxMessage); len(d) == 2 {
		if n := int(binary.BigEndian.Uint16(d)); n >= 576 {
			return n
		}
	}
	return 576
}

// Broadcast reports whether the client asked for broadcast replies.
func (p *Packet) Broadcast() bool { return p.Flags&FlagBroadcast != 0 }

// Reply starts a reply to a client's message: the same transaction,
// hardware address, flags and relay address, and the message type.
func (p *Packet) Reply(t MessageType) *Packet {
	r := &Packet{
		Op: BootReply, HType: p.HType, HLen: p.HLen, XID: p.XID, Flags: p.Flags,
		CIAddr: zero4, YIAddr: zero4, SIAddr: zero4, GIAddr: p.GIAddr, CHAddr: p.CHAddr,
	}
	if !r.GIAddr.IsValid() {
		r.GIAddr = zero4
	}
	r.Set(OptMessageType, []byte{byte(t)})
	return r
}

// NewRequest starts a client message (for tests and the CLI's probe).
func NewRequest(t MessageType, mac net.HardwareAddr, xid uint32) *Packet {
	p := &Packet{Op: BootRequest, HType: 1, HLen: byte(len(mac)), XID: xid, CIAddr: zero4, YIAddr: zero4, SIAddr: zero4, GIAddr: zero4}
	copy(p.CHAddr[:], mac)
	p.Set(OptMessageType, []byte{byte(t)})
	return p
}

// IsZero reports whether a is unset or 0.0.0.0.
func IsZero(a netip.Addr) bool { return !a.IsValid() || a == zero4 }

// U32 encodes a 32-bit option value.
func U32(v uint32) []byte { return binary.BigEndian.AppendUint32(nil, v) }

// IP encodes an address option value.
func IP(a netip.Addr) []byte {
	v := a.As4()
	return v[:]
}