package dhcp

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"strings"
)

const (
	MAGIC_COOKIE               uint32 = 0x63825363
	OPT_CODE_SERVER_IDENTIFIER uint8  = 54
	OPT_CODE_SUBNET_MASK       uint8  = 1
	OPT_CODE_ROUTER            uint8  = 3
	OPT_CODE_DNS               uint8  = 6
)

type dhcpDiscover struct {
	MsgType             byte
	HwType              byte
	HwAddrLen           byte
	Hops                byte
	Xid                 uint32
	Secs                uint16
	BootpFlags          uint16
	ClientIp            uint32
	YourIp              uint32
	NextServerIp        uint32
	RelayAgentIp        uint32
	ClientMacAddr       [6]byte
	ClientHwAddrPadding [10]byte
	ServerHostName      [64]byte
	BootFileName        [128]byte
	MagicCookie         uint32
	Mtype               byte
	MtypeLen            byte
	MtypeVal            byte
	Opts                [200]byte
	End                 byte
	Padding             [16]byte
}

func doDiscover(address string) (map[string]string, error) {

	localAddr, err := net.ResolveUDPAddr("udp", address+":68")
	if err != nil {
		return nil, err
	}
	conn, err := net.ListenUDP("udp", localAddr)
	if err != nil {
		return nil, err
	}
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	sendDiscover(conn) // broadcast 255.255.255.255:67

	return readOffer(conn) // unicast

}

func sendDiscover(conn *net.UDPConn) error {

	server := fmt.Sprintf("%s:%d", "255.255.255.255", 67)
	serverAddr, err := net.ResolveUDPAddr("udp", server)
	if err != nil {
		return err
	}

	dhcp := dhcpDiscover{
		MsgType:      0x01,
		HwType:       0x01,
		HwAddrLen:    0x06,
		Hops:         0x00,
		Xid:          0x00000000,
		Secs:         0x0000,
		ClientIp:     0x00000000,
		YourIp:       0x00000000,
		NextServerIp: 0x00000000,
		RelayAgentIp: 0x00000000,
		MagicCookie:  MAGIC_COOKIE,
		Mtype:        0x35,
		MtypeLen:     0x01,
		MtypeVal:     0x01,
		End:          0xff,
	}

	var flag uint16 = 0
	dhcp.BootpFlags = ^flag // flag = unicast , ^flag = broadcast

	writer := new(bytes.Buffer)
	if err := binary.Write(writer, binary.BigEndian, dhcp); err != nil {
		return err
	}
	_, err = conn.WriteToUDP(writer.Bytes(), serverAddr)
	if err != nil {
		return err
	}

	return nil
}

func readOffer(conn *net.UDPConn) (map[string]string, error) {
	offerOptions := make(map[string]string)

	buf := make([]byte, 1024)
	_, _, err := conn.ReadFromUDP(buf)
	if err != nil {
		return nil, err
	}

	offer := dhcpDiscover{}

	reader := new(bytes.Buffer)
	reader.Write(buf)
	if err := binary.Read(reader, binary.BigEndian, &offer); err != nil {
		return nil, err
	}
	if offer.MagicCookie != MAGIC_COOKIE {
		return nil, err
	}

	//option searching
	r := new(bytes.Buffer)
	r.Write(offer.Opts[:])

	for i := 0; i < r.Len(); i++ {
		v := r.Next(1)[0]

		if v == OPT_CODE_SUBNET_MASK && r.Next(1)[0] == 4 {
			ipStr := byteToIPString(r.Next(4))
			offerOptions["SUBNETMASK"] = ipStr
		}

		if v == OPT_CODE_ROUTER && r.Next(1)[0] == 4 {
			ipStr := byteToIPString(r.Next(4))
			offerOptions["ROUTER"] = ipStr
		}

		if v == OPT_CODE_DNS {
			len := r.Next(1)[0]
			var dns []string
			var ipStr string
			ipStr = byteToIPString(r.Next(4))
			dns = append(dns, ipStr)
			if len == 8 {
				ipStr = byteToIPString(r.Next(4))
				dns = append(dns, ipStr)
			}
			offerOptions["DNS"] = strings.Join(dns, ",")
		}

		if v == OPT_CODE_SERVER_IDENTIFIER && r.Next(1)[0] == 4 {
			ipStr := byteToIPString(r.Next(4))
			offerOptions["DHCP_SERVER"] = ipStr
		}
	}

	return offerOptions, nil
}

func byteToIPString(b []byte) string {
	var ipStr string
	len := len(b)
	for i := 0; i < len; i++ {
		v := b[i]
		ipStr += fmt.Sprintf("%d", v)
		if i < len-1 {
			ipStr += "."
		}
	}

	return ipStr
}