package node

import (
	"encoding/binary"
	"net"
)

type LuxDnsServer struct {
	node     *LuxNode
	stopChan chan bool
}

func NewLuxDnsServer(node *LuxNode) *LuxDnsServer {
	return &LuxDnsServer{
		node:     node,
		stopChan: make(chan bool),
	}
}

var NO = binary.BigEndian

func (sv *LuxDnsServer) HandleRequest(req []byte) []byte {
	id := NO.Uint16(req[0:2])
	flags := NO.Uint16(req[2:4])
	qdcount := NO.Uint16(req[4:6])
	ancount := NO.Uint16(req[6:8])
	nscount := NO.Uint16(req[8:10])
	arcount := NO.Uint16(req[10:12])
	data := req[12:]

	qr := flags >> 15
	opcode := (flags >> 11) & 0b1111
	var rcode uint16

	if qr != 0 || opcode != 0 {
		// throw not supported
		goto reply
	}

	qr = 1
	rcode = 4
reply:
	flags |= (qr << 15) | rcode

	hdr := make([]byte, 12)
	NO.PutUint16(hdr[0:2], id)
	NO.PutUint16(hdr[2:4], flags)
	NO.PutUint16(hdr[4:6], qdcount)
	NO.PutUint16(hdr[6:8], ancount)
	NO.PutUint16(hdr[8:10], nscount)
	NO.PutUint16(hdr[10:12], arcount)

	return append(hdr, data...)
}

func (sv *LuxDnsServer) CreateFrontend(udpListen string) error {
	listenAddr, err := net.ResolveUDPAddr("udp", udpListen)
	if err != nil {
		return err
	}

	conn, err := net.ListenUDP("udp", listenAddr)
	if err != nil {
		return err
	}

	go func(sv *LuxDnsServer) {
		defer conn.Close()

		buf := make([]byte, 512)

		for {
			select {
			case <-sv.stopChan:
				return
			default:
				n, addr, err := conn.ReadFromUDP(buf)
				if err != nil {
					log.Debugf("failed to recv dns udp: %v\n", err)
				}

				res := sv.HandleRequest(buf[:n])
				if len(res) > 0 {
					_, err := conn.WriteToUDP(res, addr)
					if err != nil {
						log.Debugf("failed to send dns reply: %v\n", err)
					}
				}
			}
		}
	}(sv)

	return nil
}

func (sv *LuxDnsServer) Stop() {
	sv.stopChan <- true
}