package nettools

import (
	"crypto/tls"
	"errors"
	"fmt"
	"github.com/miekg/dns"
	"golang.org/x/net/context"
	"net"
	"strconv"
	"strings"
	"time"
)

type DnsPingResult struct {
	Time int
	Err  error
	IP   net.IP
}

func (dnsR *DnsPingResult) Result() int {
	return dnsR.Time
}

func (dnsR *DnsPingResult) Error() error {
	return dnsR.Err
}

func (dnsR *DnsPingResult) String() string {
	if dnsR.Err != nil {
		return fmt.Sprintf("%s", dnsR.Err)
	} else {
		return fmt.Sprintf("%s: time=%d ms", dnsR.IP.String(), dnsR.Time)
	}
}

type DnsPing struct {
	host    string
	Port    uint16
	Timeout time.Duration

	// udp, tcp, tcp-tls，默认 udp
	Net string

	// A, AAAA, NS, ...，默认 NS
	Type string

	// 查询域名，默认 .
	Domain string

	// Net 为 tcp-tls 时，是否跳过证书验证
	Insecure bool

	ip net.IP
}

func (dnsC *DnsPing) SetHost(host string) {
	dnsC.host = host
	dnsC.ip = net.ParseIP(host)
}

func (dnsC *DnsPing) Host() string {
	return dnsC.host
}

func (dnsC *DnsPing) Ping() IPingResult {
	return dnsC.PingContext(context.Background())
}

func (dnsC *DnsPing) PingContext(ctx context.Context) IPingResult {
	ip := cloneIP(dnsC.ip)
	if ip == nil {
		var err error
		ip, err = LookupFunc(dnsC.host)
		if err != nil {
			return &DnsPingResult{0, err, nil}
		}
	}

	msg := &dns.Msg{}
	qtype, ok := dns.StringToType[dnsC.Type]
	if !ok {
		return &DnsPingResult{0, errors.New("unknown type"), nil}
	}
	if !strings.HasSuffix(dnsC.Domain, ".") {
		dnsC.Domain += "."
	}
	msg.SetQuestion(dnsC.Domain, qtype)
	msg.MsgHdr.RecursionDesired = true

	client := &dns.Client{}
	client.Net = dnsC.Net
	client.Timeout = dnsC.Timeout
	client.TLSConfig = &tls.Config{
		ServerName:         dnsC.host,
		InsecureSkipVerify: dnsC.Insecure,
	}

	t0 := time.Now()
	r, _, err := client.ExchangeContext(ctx, msg, net.JoinHostPort(ip.String(), strconv.Itoa(int(dnsC.Port))))
	if err != nil {
		return &DnsPingResult{0, err, nil}
	}
	if r == nil || r.Response == false || r.Opcode != dns.OpcodeQuery {
		return &DnsPingResult{0, errors.New("response error"), nil}
	}
	return &DnsPingResult{int(time.Now().Sub(t0).Milliseconds()), nil, ip}
}

func NewDnsPing(host string, timeout time.Duration) *DnsPing {
	return &DnsPing{
		host:     host,
		Port:     53,
		Timeout:  timeout,
		Net:      "udp",
		Type:     "NS",
		Domain:   ".",
		Insecure: false,
		ip:       net.ParseIP(host),
	}
}

var (
	_ IPing       = (*DnsPing)(nil)
	_ IPingResult = (*DnsPingResult)(nil)
)