package dnstt

import (
	"fmt"
	"net/netip"
	"strconv"

	geoip2 "github.com/oschwald/geoip2-golang/v2"
)

const (
	// UnknownCountry is used when country lookup is enabled but no country is found.
	UnknownCountry = "ZZ"
	// UnknownASN is used when ASN lookup is enabled but no ASN is found.
	UnknownASN = "0"
)

// GeoLabels are optional labels derived from a resolver IP address.
type GeoLabels struct {
	Country string
	ASN     string
}

// GeoResolver resolves optional GeoIP labels for resolver IP addresses.
type GeoResolver interface {
	Lookup(netip.Addr) GeoLabels
	LabelNames() []string
}

// GeoIPResolver uses optional MaxMind Country and ASN databases.
type GeoIPResolver struct {
	countryDB *geoip2.Reader
	asnDB     *geoip2.Reader
}

// OpenGeoIPResolver opens optional MaxMind Country and ASN databases.
func OpenGeoIPResolver(countryDatabase string, asnDatabase string) (*GeoIPResolver, error) {
	resolver := &GeoIPResolver{}
	if countryDatabase != "" {
		db, err := geoip2.Open(countryDatabase)
		if err != nil {
			return nil, fmt.Errorf("open country database: %w", err)
		}
		resolver.countryDB = db
	}
	if asnDatabase != "" {
		db, err := geoip2.Open(asnDatabase)
		if err != nil {
			resolver.Close()
			return nil, fmt.Errorf("open ASN database: %w", err)
		}
		resolver.asnDB = db
	}
	return resolver, nil
}

// Close closes any open MaxMind databases.
func (r *GeoIPResolver) Close() {
	if r.countryDB != nil {
		r.countryDB.Close()
	}
	if r.asnDB != nil {
		r.asnDB.Close()
	}
}

// LabelNames returns the optional Prometheus labels enabled by configured databases.
func (r *GeoIPResolver) LabelNames() []string {
	var labels []string
	if r.countryDB != nil {
		labels = append(labels, "country")
	}
	if r.asnDB != nil {
		labels = append(labels, "asn")
	}
	return labels
}

// Lookup returns optional GeoIP labels for a resolver IP address.
func (r *GeoIPResolver) Lookup(addr netip.Addr) GeoLabels {
	labels := GeoLabels{}
	if !addr.IsValid() {
		if r.countryDB != nil {
			labels.Country = UnknownCountry
		}
		if r.asnDB != nil {
			labels.ASN = UnknownASN
		}
		return labels
	}

	if r.countryDB != nil {
		labels.Country = UnknownCountry
		record, err := r.countryDB.Country(addr)
		if err == nil && record.HasData() && record.Country.ISOCode != "" {
			labels.Country = record.Country.ISOCode
		}
	}
	if r.asnDB != nil {
		labels.ASN = UnknownASN
		record, err := r.asnDB.ASN(addr)
		if err == nil && record.HasData() && record.AutonomousSystemNumber != 0 {
			labels.ASN = strconv.FormatUint(uint64(record.AutonomousSystemNumber), 10)
		}
	}

	return labels
}