package network

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"strings"
	"time"
)

// The record types a look asks for.
const (
	TypeA    uint16 = 1
	TypeAAAA uint16 = 28
)

// The answers a resolver gives that a look tells apart.
const (
	RcodeOK       = 0
	RcodeServFail = 2
	RcodeNXDomain = 3
	RcodeRefused  = 5
)

// Answer is what one resolver said to one question: its code, how many records of the type asked it
// gave, and how long it took. A question with no answer at all is an error, never an Answer.
type Answer struct {
	Rcode   int
	Records int
	Took    time.Duration
}

// Ask asks one resolver one question over UDP, and reads its code and how many records of the type
// asked it answered with. **It tells "no such name" from "no record of that type"** — the C library's
// lookup does not, and that difference is issue 262: an Alpine container failed a mesh name because a
// resolver said NXDOMAIN for its IPv6 address where it should have said there was none.
func Ask(ctx context.Context, server, name string, qtype uint16, timeout time.Duration) (Answer, error) {
	start := time.Now()
	query, id, err := question(name, qtype)
	if err != nil {
		return Answer{}, err
	}
	if net.ParseIP(server) != nil {
		server = net.JoinHostPort(server, "53")
	}
	d := net.Dialer{Timeout: timeout}
	conn, err := d.DialContext(ctx, "udp", server)
	if err != nil {
		return Answer{}, err
	}
	defer conn.Close()
	_ = conn.SetDeadline(start.Add(timeout))
	if _, err := conn.Write(query); err != nil {
		return Answer{}, err
	}
	buf := make([]byte, 1500)
	for {
		n, err := conn.Read(buf)
		if err != nil {
			var ne net.Error
			if errors.As(err, &ne) && ne.Timeout() {
				return Answer{}, fmt.Errorf("no answer within %s", timeout)
			}
			return Answer{}, err
		}
		a, ours, err := parse(buf[:n], id, qtype)
		if !ours {
			continue // a late answer to somebody else's question on this port
		}
		if err != nil {
			return Answer{}, err
		}
		a.Took = time.Since(start)
		return a, nil
	}
}

// question is one query: a header asking for recursion, and the name and type.
func question(name string, qtype uint16) ([]byte, uint16, error) {
	var idb [2]byte
	if _, err := rand.Read(idb[:]); err != nil {
		return nil, 0, err
	}
	id := binary.BigEndian.Uint16(idb[:])
	msg := make([]byte, 12, 64)
	binary.BigEndian.PutUint16(msg[0:], id)
	msg[2] = 0x01 // recursion desired
	binary.BigEndian.PutUint16(msg[4:], 1)
	for _, label := range strings.Split(strings.TrimSuffix(name, "."), ".") {
		if label == "" || len(label) > 63 {
			return nil, 0, fmt.Errorf("%q is not a name that can be asked", name)
		}
		msg = append(msg, byte(len(label)))
		msg = append(msg, label...)
	}
	msg = append(msg, 0, byte(qtype>>8), byte(qtype), 0, 1)
	return msg, id, nil
}

// parse reads an answer: false when it is not the answer to this question.
func parse(msg []byte, id, qtype uint16) (Answer, bool, error) {
	if len(msg) < 12 || binary.BigEndian.Uint16(msg[0:]) != id || msg[2]&0x80 == 0 {
		return Answer{}, false, nil
	}
	a := Answer{Rcode: int(msg[3] & 0x0f)}
	qd, an := int(binary.BigEndian.Uint16(msg[4:])), int(binary.BigEndian.Uint16(msg[6:]))
	at := 12
	for i := 0; i < qd; i++ {
		var err error
		if at, err = skipName(msg, at); err != nil {
			return a, true, err
		}
		at += 4
	}
	for i := 0; i < an; i++ {
		var err error
		if at, err = skipName(msg, at); err != nil {
			return a, true, err
		}
		if at+10 > len(msg) {
			return a, true, errors.New("the answer is cut short")
		}
		typ := binary.BigEndian.Uint16(msg[at:])
		length := int(binary.BigEndian.Uint16(msg[at+8:]))
		at += 10 + length
		if at > len(msg) {
			return a, true, errors.New("the answer is cut short")
		}
		if typ == qtype {
			a.Records++
		}
	}
	return a, true, nil
}

// skipName steps over a name, compressed or not.
func skipName(msg []byte, at int) (int, error) {
	for {
		if at >= len(msg) {
			return 0, errors.New("the answer is cut short")
		}
		l := int(msg[at])
		switch {
		case l == 0:
			return at + 1, nil
		case l&0xc0 == 0xc0:
			return at + 2, nil
		default:
			at += 1 + l
		}
	}
}