package doh

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

const (
	MsgDnsNotReachable = "DNS_NOT_REACHABLE"
	MsgDnsProbeFailed  = "DNS_PROBE_FAILED"
)

type googleDohResponse struct {
	Status int   `json:"Status"`
	Answer []any `json:"Answer"`
}

type probeResult struct {
	reachable bool
	msg       string
	err       error
}

var dohEndpoints = map[string]string{
	"google":     "https://dns.google/resolve?name=%s&type=A",
	"cloudflare": "https://cloudflare-dns.com/dns-query?name=%s&type=A",
}

func CheckDnsPropagation(ctx context.Context, dns string) (bool, string) {
	results := make(chan probeResult, len(dohEndpoints))

	for _, url := range dohEndpoints {
		go func(u, d string) {
			res, msg, err := fetchDoh(ctx, fmt.Sprintf(u, d))
			results <- probeResult{reachable: res, msg: msg, err: err}
		}(url, dns)
	}

	var validResults []probeResult
	var probeErrors []error

	for i := 0; i < len(dohEndpoints); i++ {
		res := <-results
		if res.err != nil {
			probeErrors = append(probeErrors, res.err)
		} else {
			validResults = append(validResults, res)
		}
	}

	if len(probeErrors) > 0 {
		for _, probeErr := range probeErrors {
			tilog.Warn(ctx, "DNS over HTTPS probe error", zap.Error(probeErr))
		}
	}

	if len(validResults) > 0 {
		for _, v := range validResults {
			if v.reachable {
				return true, ""
			}
		}
		return false, MsgDnsNotReachable
	}
	return false, MsgDnsProbeFailed
}

func fetchDoh(ctx context.Context, apiURL string) (bool, string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, apiURL, http.NoBody)
	if err != nil {
		return false, "", err
	}
	req.Header.Set("Accept", "application/dns-json")

	client := &http.Client{Timeout: 5 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return false, "", err
	}
	defer func() {
		if resp != nil && resp.Body != nil {
			_ = resp.Body.Close()
		}
	}()

	if resp.StatusCode != http.StatusOK {
		return false, "", fmt.Errorf("http status %d", resp.StatusCode)
	}

	var dohResp googleDohResponse
	if err := json.NewDecoder(resp.Body).Decode(&dohResp); err != nil {
		return false, "", err
	}

	switch dohResp.Status {
	case 0:
		if len(dohResp.Answer) > 0 {
			return true, "", nil
		}
		return false, MsgDnsNotReachable, nil
	case 3:
		return false, MsgDnsNotReachable, nil
	default:
		return false, "", fmt.Errorf("dns status %d", dohResp.Status)
	}
}