package ws

import (
	"git.smesh.lol/moxie/pkg/mxutil"
	"fmt"
	"os"
	"sync"
	"syscall"
	"time"
)

const dnsTTL = 24 * time.Hour

var dnsAddr [4]byte
var dnsAddrInit bool

func initDNSAddr() {
	if dnsAddrInit {
		return
	}
	dnsAddrInit = true
	dnsAddr = [4]byte{8, 8, 8, 8}
	data, err := os.ReadFile("/etc/resolv.conf")
	if err != nil {
		return
	}
	i := 0
	for i < len(data) {
		lineStart := i
		for i < len(data) && data[i] != '\n' {
			i++
		}
		line := data[lineStart:i]
		if i < len(data) {
			i++
		}
		if len(line) > 11 && string(line[:11]) == "nameserver " {
			ns := line[11:]
			for len(ns) > 0 && (ns[len(ns)-1] == ' ' || ns[len(ns)-1] == '\t' || ns[len(ns)-1] == '\r') {
				ns = ns[:len(ns)-1]
			}
			if ip := parseIPv4(ns); ip != nil {
				dnsAddr = [4]byte{ip[0], ip[1], ip[2], ip[3]}
				return
			}
		}
	}
}

func parseIPv4(s []byte) (buf []byte) {
	var parts [4]byte
	p := 0
	v := 0
	for i := 0; i <= len(s); i++ {
		if i == len(s) || s[i] == '.' {
			if v > 255 || p > 3 {
				return nil
			}
			parts[p] = byte(v)
			p++
			v = 0
		} else if s[i] >= '0' && s[i] <= '9' {
			v = v*10 + int32(s[i]-'0')
		} else {
			return nil
		}
	}
	if p != 4 {
		return nil
	}
	return parts[:]
}

type dnsEntry struct {
	ip  string
	exp time.Time
}

// dnsCacheState owns the resolver cache. A cache written after init must live
// behind a self-mutating type: the rule forbids stores to a package global
// outside init, and a method on the owner is the sanctioned place for them.
type dnsCacheState struct {
	mu sync.Mutex
	m  map[string]dnsEntry
}

func (c *dnsCacheState) Get(host string) (e dnsEntry, ok bool) {
	c.mu.Lock()
	e, ok = c.m[host]
	c.mu.Unlock()
	return
}

func (c *dnsCacheState) Set(host string, e dnsEntry) {
	c.mu.Lock()
	c.m[host] = e
	c.mu.Unlock()
}

var dnsCache dnsCacheState

func init() {
	dnsCache.m = map[string]dnsEntry{}
}

// ResolveHost is the exported entrypoint for the cached DNS resolver.
// Used by other packages that need to dial a hostname without going
// through net/http (which is broken in Moxie).
func ResolveHost(host string) (addr string, err error) { return resolveHost(host) }

// resolveHost returns a cached IP for the hostname, or resolves via raw UDP DNS.
func resolveHost(host string) (addr string, derr error) {
	if e, ok := dnsCache.Get(host); ok && time.Now().Before(e.exp) {
		return e.ip, nil
	}

	ip, err := dnsLookup(host)
	if err != nil {
		return "", err
	}

	dnsCache.Set(host, dnsEntry{ip: ip, exp: time.Now().Add(dnsTTL)})

	return ip, nil
}

// dnsLookup sends a minimal DNS A-record query via raw syscalls.
func dnsLookup(host string) (addr string, derr error) {
	fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_DGRAM, 0)
	if err != nil {
		return "", fmt.Errorf("dns: socket: %w", err)
	}
	defer syscall.Close(fd)

	tv := syscall.Timeval{Sec: 5}
	syscall.SetsockoptTimeval(fd, syscall.SOL_SOCKET, syscall.SO_RCVTIMEO, &tv)

	initDNSAddr()
	sa := &syscall.SockaddrInet4{Port: 53, Addr: dnsAddr}
	if err = syscall.Connect(fd, sa); err != nil {
		return "", fmt.Errorf("dns: connect: %w", err)
	}

	// Build DNS A-record query.
	var pkt []byte
	pkt = push(pkt, 0xAB, 0xCD)                         // ID
	pkt = push(pkt, 0x01, 0x00)                         // flags: recursion desired
	pkt = push(pkt, 0x00, 0x01)                         // 1 question
	pkt = push(pkt, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00) // 0 answer/auth/additional

	// Encode QNAME.
	j := 0
	for j < len(host) {
		dot := j
		for dot < len(host) && host[dot] != '.' {
			dot++
		}
		pkt = mxutil.Ensure(pkt, 1)
		pkt = push(pkt, byte(dot-j))
		pkt = pkt | host[j:dot]
		j = dot + 1
	}
	pkt = push(pkt, 0x00)       // root
	pkt = push(pkt, 0x00, 0x01) // QTYPE A
	pkt = push(pkt, 0x00, 0x01) // QCLASS IN

	if err = syscall.Sendto(fd, pkt, 0, sa); err != nil {
		return "", fmt.Errorf("dns: send: %w", err)
	}

	buf := []byte{:512}
	n, _, err := syscall.Recvfrom(fd, buf, 0)
	if err != nil {
		return "", fmt.Errorf("dns: recv: %w", err)
	}
	if n < 12 {
		return "", fmt.Errorf("dns: response too short")
	}

	anCount := int32(buf[6])<<8 | int32(buf[7])
	if anCount == 0 {
		return "", fmt.Errorf("dns: no answers for %s", host)
	}

	// Skip question section.
	pos := 12
	for pos < n {
		if buf[pos] == 0 {
			pos++
			break
		}
		if buf[pos]&0xC0 == 0xC0 {
			pos += 2
			break
		}
		pos += int32(buf[pos]) + 1
	}
	pos += 4 // QTYPE + QCLASS

	// Find first A record.
	for a := 0; a < anCount && pos+10 < n; a++ {
		if buf[pos]&0xC0 == 0xC0 {
			pos += 2
		} else {
			for pos < n && buf[pos] != 0 {
				pos += int32(buf[pos]) + 1
			}
			pos++
		}
		if pos+10 > n {
			break
		}
		rtype := int32(buf[pos])<<8 | int32(buf[pos+1])
		rdlen := int32(buf[pos+8])<<8 | int32(buf[pos+9])
		pos += 10
		if rtype == 1 && rdlen == 4 && pos+4 <= n {
			// Copy Sprintf result to avoid buffer reuse corruption.
			// The bytes are widened to int32: Moxie's fmt type-switch
			// matches uint8 but not the named type byte, so %d on a
			// byte prints "?" and the resolver would return "?.?.?.?"
			// as a successful address.
			s := fmt.Sprintf("%d.%d.%d.%d", int32(buf[pos]), int32(buf[pos+1]), int32(buf[pos+2]), int32(buf[pos+3]))
			ip := []byte{:len(s)}
			copy(ip, s)
			return string(ip), nil
		}
		pos += rdlen
	}

	return "", fmt.Errorf("dns: no A record for %s", host)
}