package dnscfg

import (
	"context"
	"fmt"
	"net"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	BenchmarkProfileQuick = "quick"
	BenchmarkProfileLoad  = "load"

	BenchmarkErrorNXDomain  = "nxdomain"
	BenchmarkErrorTimeout   = "timeout"
	BenchmarkErrorTemporary = "temporary"
	BenchmarkErrorOther     = "other"
)

var BenchmarkDefaultDomains = []string{
	"cloudflare.com",
	"google.com",
	"telegram.org",
	"github.com",
	"youtube.com",
	"twitter.com",
}

type BenchmarkOptions struct {
	Profile            string
	LoadWorkers        int
	Rounds             int
	SyntheticPerDomain int
}

type BenchmarkResult struct {
	Upstream  string
	Attempts  int
	OK        int
	Fail      int
	NXDomain  int
	Timeout   int
	Temporary int
	Other     int
	AvgMS     int
	P95MS     int
	Score     float64
	Color     string
}

func NormalizeBenchmarkProfile(raw string) string {
	switch strings.ToLower(strings.TrimSpace(raw)) {
	case "", BenchmarkProfileLoad:
		return BenchmarkProfileLoad
	case BenchmarkProfileQuick:
		return BenchmarkProfileQuick
	default:
		return BenchmarkProfileLoad
	}
}

func MakeDNSBenchmarkOptions(profile string, concurrency int) BenchmarkOptions {
	if concurrency < 1 {
		concurrency = 1
	}
	if profile == BenchmarkProfileQuick {
		return BenchmarkOptions{
			Profile:            BenchmarkProfileQuick,
			LoadWorkers:        1,
			Rounds:             1,
			SyntheticPerDomain: 0,
		}
	}
	workers := concurrency * 2
	if workers < 4 {
		workers = 4
	}
	if workers > 16 {
		workers = 16
	}
	return BenchmarkOptions{
		Profile:            BenchmarkProfileLoad,
		LoadWorkers:        workers,
		Rounds:             3,
		SyntheticPerDomain: 2,
	}
}

func NormalizeBenchmarkUpstreamStrings(in []string, normalizeUpstream func(string, string) string) []string {
	out := make([]string, 0, len(in))
	seen := map[string]struct{}{}
	for _, raw := range in {
		n := strings.TrimSpace(raw)
		if normalizeUpstream != nil {
			n = normalizeUpstream(n, "53")
		}
		if n == "" {
			continue
		}
		if _, ok := seen[n]; ok {
			continue
		}
		seen[n] = struct{}{}
		out = append(out, n)
	}
	return out
}

func NormalizeBenchmarkDomains(in []string) []string {
	if len(in) == 0 {
		return nil
	}
	out := make([]string, 0, len(in))
	seen := map[string]struct{}{}
	for _, raw := range in {
		d := strings.TrimSuffix(strings.ToLower(strings.TrimSpace(raw)), ".")
		if d == "" || strings.HasPrefix(d, "#") {
			continue
		}
		if _, ok := seen[d]; ok {
			continue
		}
		seen[d] = struct{}{}
		out = append(out, d)
	}
	if len(out) > 100 {
		out = out[:100]
	}
	return out
}

func BenchmarkDNSUpstream(
	upstream string,
	domains []string,
	timeout time.Duration,
	attempts int,
	opts BenchmarkOptions,
	lookupAOnce func(host, upstream string, timeout time.Duration) ([]string, error),
	classifyErr func(error) string,
) BenchmarkResult {
	res := BenchmarkResult{Upstream: upstream}
	probes := BuildBenchmarkProbeHosts(domains, attempts, opts)
	if len(probes) == 0 {
		return res
	}
	durations := make([]int, 0, len(probes))
	var mu sync.Mutex
	jobs := make(chan string, len(probes))
	workers := opts.LoadWorkers
	if workers < 1 {
		workers = 1
	}
	if workers > len(probes) {
		workers = len(probes)
	}
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for host := range jobs {
				start := time.Now()
				_, err := lookupAOnce(host, upstream, timeout)
				elapsed := int(time.Since(start).Milliseconds())
				if elapsed < 1 {
					elapsed = 1
				}
				mu.Lock()
				res.Attempts++
				durations = append(durations, elapsed)
				if err != nil {
					res.Fail++
					switch strings.ToLower(strings.TrimSpace(classifyErr(err))) {
					case BenchmarkErrorNXDomain:
						res.NXDomain++
					case BenchmarkErrorTimeout:
						res.Timeout++
					case BenchmarkErrorTemporary:
						res.Temporary++
					default:
						res.Other++
					}
				} else {
					res.OK++
				}
				mu.Unlock()
			}
		}()
	}
	for _, host := range probes {
		jobs <- host
	}
	close(jobs)
	wg.Wait()

	if len(durations) > 0 {
		sort.Ints(durations)
		sum := 0
		for _, d := range durations {
			sum += d
		}
		res.AvgMS = sum / len(durations)
		idx := int(float64(len(durations)-1) * 0.95)
		if idx < 0 {
			idx = 0
		}
		res.P95MS = durations[idx]
	}

	total := res.Attempts
	if total > 0 {
		okRate := float64(res.OK) / float64(total)
		answeredRate := float64(res.OK+res.NXDomain+res.Temporary+res.Other) / float64(total)
		timeoutRate := float64(res.Timeout) / float64(total)
		temporaryRate := float64(res.Temporary) / float64(total)
		otherRate := float64(res.Other) / float64(total)
		avg := float64(res.AvgMS)
		if avg <= 0 {
			avg = float64(timeout.Milliseconds())
		}
		p95 := float64(res.P95MS)
		if p95 <= 0 {
			p95 = avg
		}
		res.Score = answeredRate*100.0 + okRate*15.0 - timeoutRate*120.0 - temporaryRate*35.0 - otherRate*20.0 - (avg / 25.0) - (p95 / 45.0)
	}

	timeoutRate := 0.0
	answeredRate := 0.0
	if res.Attempts > 0 {
		timeoutRate = float64(res.Timeout) / float64(res.Attempts)
		answeredRate = float64(res.OK+res.NXDomain+res.Temporary+res.Other) / float64(res.Attempts)
	}
	switch {
	case answeredRate < 0.85 || timeoutRate >= 0.10 || res.P95MS > 1800:
		res.Color = "red"
	case answeredRate >= 0.97 && timeoutRate <= 0.02 && res.P95MS <= 700:
		res.Color = "green"
	default:
		res.Color = "yellow"
	}

	return res
}

func BuildBenchmarkProbeHosts(domains []string, attempts int, opts BenchmarkOptions) []string {
	if len(domains) == 0 {
		return nil
	}
	if attempts < 1 {
		attempts = 1
	}
	rounds := opts.Rounds
	if rounds < 1 {
		rounds = 1
	}
	synth := opts.SyntheticPerDomain
	if synth < 0 {
		synth = 0
	}
	out := make([]string, 0, len(domains)*attempts*rounds*(1+synth))
	for round := 0; round < rounds; round++ {
		for _, host := range domains {
			for i := 0; i < attempts; i++ {
				out = append(out, host)
			}
			for n := 0; n < synth; n++ {
				out = append(out, fmt.Sprintf("svpn-bench-%d-%d.%s", round+1, n+1, host))
			}
		}
	}
	if len(out) > 10000 {
		out = out[:10000]
	}
	return out
}

func DNSLookupAOnce(
	host string,
	upstream string,
	timeout time.Duration,
	splitDNS func(string) (string, string),
	isPrivateIPv4 func(string) bool,
) ([]string, error) {
	if splitDNS == nil {
		return nil, fmt.Errorf("splitDNS callback is nil")
	}
	server, port := splitDNS(upstream)
	if server == "" {
		return nil, fmt.Errorf("upstream empty")
	}
	if port == "" {
		port = "53"
	}
	addr := net.JoinHostPort(server, port)
	resolver := &net.Resolver{
		PreferGo: true,
		Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
			d := net.Dialer{}
			return d.DialContext(ctx, "udp", addr)
		},
	}

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	records, err := resolver.LookupHost(ctx, host)
	cancel()
	if err != nil {
		return nil, err
	}
	seen := map[string]struct{}{}
	out := make([]string, 0, len(records))
	for _, ip := range records {
		if isPrivateIPv4 != nil && isPrivateIPv4(ip) {
			continue
		}
		if _, ok := seen[ip]; ok {
			continue
		}
		seen[ip] = struct{}{}
		out = append(out, ip)
	}
	if len(out) == 0 {
		return nil, fmt.Errorf("no public ips")
	}
	return out, nil
}

func BenchmarkTopN(results []BenchmarkResult, n int, fallback []string) []string {
	out := make([]string, 0, n)
	for _, item := range results {
		if item.OK <= 0 {
			continue
		}
		out = append(out, item.Upstream)
		if len(out) >= n {
			return out
		}
	}
	for _, item := range fallback {
		if len(out) >= n {
			break
		}
		dup := false
		for _, e := range out {
			if e == item {
				dup = true
				break
			}
		}
		if !dup {
			out = append(out, item)
		}
	}
	return out
}