package dnstt

import (
	"net/netip"
	"strings"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus/testutil"
)

func TestExporterCollectsAggregateDNSTTMetrics(t *testing.T) {
	now := time.Unix(1000, 0)
	c := NewCollector([]string{"tunnel.example.com"}, WithNow(func() time.Time { return now }))
	c.RecordQuery("tunnel.example.com", "client-a", 100)
	c.RecordQuery("tunnel.example.com", "client-b", 300)
	c.RecordResponse("tunnel.example.com", 250)

	expected := `
# HELP dnstt_active_clients Number of DNSTT client sessions observed within the active timeout window.
# TYPE dnstt_active_clients gauge
dnstt_active_clients{domain="tunnel.example.com"} 2
# HELP dnstt_bytes_in_total Total bytes observed in DNSTT DNS queries.
# TYPE dnstt_bytes_in_total counter
dnstt_bytes_in_total{domain="tunnel.example.com"} 400
# HELP dnstt_bytes_out_total Total bytes observed in DNSTT DNS responses.
# TYPE dnstt_bytes_out_total counter
dnstt_bytes_out_total{domain="tunnel.example.com"} 250
# HELP dnstt_peak_clients Maximum concurrent active DNSTT client sessions observed.
# TYPE dnstt_peak_clients gauge
dnstt_peak_clients{domain="tunnel.example.com"} 2
# HELP dnstt_queries_total Total DNSTT DNS queries observed.
# TYPE dnstt_queries_total counter
dnstt_queries_total{domain="tunnel.example.com"} 2
# HELP dnstt_sessions_total Total unique DNSTT client sessions observed.
# TYPE dnstt_sessions_total counter
dnstt_sessions_total{domain="tunnel.example.com"} 2
`

	if err := testutil.CollectAndCompare(NewExporter(c), strings.NewReader(expected)); err != nil {
		t.Fatal(err)
	}
}

func TestExporterCollectsGeoIPCountryAndASNLabels(t *testing.T) {
	now := time.Unix(1000, 0)
	resolverIP := netip.MustParseAddr("2001:db8::53")
	c := NewCollector(
		[]string{"tunnel.example.com"},
		WithNow(func() time.Time { return now }),
		WithGeoResolver(fakeGeoResolver{
			labelNames: []string{"country", "asn"},
			labels: map[netip.Addr]GeoLabels{
				resolverIP: {Country: "DE", ASN: "3320"},
			},
		}),
	)
	c.RecordQueryFrom("tunnel.example.com", "client-a", resolverIP, 100)

	expected := `
# HELP dnstt_active_clients Number of DNSTT client sessions observed within the active timeout window.
# TYPE dnstt_active_clients gauge
dnstt_active_clients{asn="3320",country="DE",domain="tunnel.example.com"} 1
# HELP dnstt_bytes_in_total Total bytes observed in DNSTT DNS queries.
# TYPE dnstt_bytes_in_total counter
dnstt_bytes_in_total{asn="3320",country="DE",domain="tunnel.example.com"} 100
# HELP dnstt_bytes_out_total Total bytes observed in DNSTT DNS responses.
# TYPE dnstt_bytes_out_total counter
dnstt_bytes_out_total{asn="3320",country="DE",domain="tunnel.example.com"} 0
# HELP dnstt_peak_clients Maximum concurrent active DNSTT client sessions observed.
# TYPE dnstt_peak_clients gauge
dnstt_peak_clients{asn="3320",country="DE",domain="tunnel.example.com"} 1
# HELP dnstt_queries_total Total DNSTT DNS queries observed.
# TYPE dnstt_queries_total counter
dnstt_queries_total{asn="3320",country="DE",domain="tunnel.example.com"} 1
# HELP dnstt_sessions_total Total unique DNSTT client sessions observed.
# TYPE dnstt_sessions_total counter
dnstt_sessions_total{asn="3320",country="DE",domain="tunnel.example.com"} 1
`

	if err := testutil.CollectAndCompare(NewExporter(c), strings.NewReader(expected)); err != nil {
		t.Fatal(err)
	}
}

func TestExporterCanCollectASNWithoutCountry(t *testing.T) {
	now := time.Unix(1000, 0)
	resolverIP := netip.MustParseAddr("192.0.2.53")
	c := NewCollector(
		[]string{"tunnel.example.com"},
		WithNow(func() time.Time { return now }),
		WithGeoResolver(fakeGeoResolver{
			labelNames: []string{"asn"},
			labels: map[netip.Addr]GeoLabels{
				resolverIP: {ASN: "15169"},
			},
		}),
	)
	c.RecordQueryFrom("tunnel.example.com", "client-a", resolverIP, 100)

	expected := `
# HELP dnstt_active_clients Number of DNSTT client sessions observed within the active timeout window.
# TYPE dnstt_active_clients gauge
dnstt_active_clients{asn="15169",domain="tunnel.example.com"} 1
# HELP dnstt_bytes_in_total Total bytes observed in DNSTT DNS queries.
# TYPE dnstt_bytes_in_total counter
dnstt_bytes_in_total{asn="15169",domain="tunnel.example.com"} 100
# HELP dnstt_bytes_out_total Total bytes observed in DNSTT DNS responses.
# TYPE dnstt_bytes_out_total counter
dnstt_bytes_out_total{asn="15169",domain="tunnel.example.com"} 0
# HELP dnstt_peak_clients Maximum concurrent active DNSTT client sessions observed.
# TYPE dnstt_peak_clients gauge
dnstt_peak_clients{asn="15169",domain="tunnel.example.com"} 1
# HELP dnstt_queries_total Total DNSTT DNS queries observed.
# TYPE dnstt_queries_total counter
dnstt_queries_total{asn="15169",domain="tunnel.example.com"} 1
# HELP dnstt_sessions_total Total unique DNSTT client sessions observed.
# TYPE dnstt_sessions_total counter
dnstt_sessions_total{asn="15169",domain="tunnel.example.com"} 1
`

	if err := testutil.CollectAndCompare(NewExporter(c), strings.NewReader(expected)); err != nil {
		t.Fatal(err)
	}
}

type fakeGeoResolver struct {
	labelNames []string
	labels     map[netip.Addr]GeoLabels
}

func (f fakeGeoResolver) Lookup(addr netip.Addr) GeoLabels {
	if labels, ok := f.labels[addr]; ok {
		return labels
	}
	return GeoLabels{Country: UnknownCountry, ASN: UnknownASN}
}

func (f fakeGeoResolver) LabelNames() []string {
	return f.labelNames
}