package geo

import (
	"context"
	"errors"
	"log"
	"math"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/oschwald/maxminddb-golang"
)

// Locator turns a client address into a point. A nil Locator, or a nil
// *MMDB, locates nothing: the IP step is optional (spec §6.2).
type Locator interface {
	Locate(ip net.IP) (Point, bool)
}

// MMDB is an operator-supplied MaxMind-format database. lookup is the
// reader's Lookup, a field so a test can prove which addresses ever reach
// it; source, build and nodes are copied from the metadata at open so a
// test fake can set them too.
type MMDB struct {
	r      *maxminddb.Reader
	lookup func(ip net.IP, result any) error
	source string
	build  time.Time
	nodes  uint64

	closeHook func() // test seam: called by Close before the mapping is released
}

// OpenMMDB opens and verifies the file. An error means "IP step off"; the
// console logs it once and starts anyway.
func OpenMMDB(path string) (*MMDB, error) {
	return OpenMMDBVerified(context.Background(), path)
}

// ErrSearchTreeTooLarge is a database whose search tree has far more paths
// than nodes: shared subtrees (a DAG) make the library's verifier, which
// enumerates every network, run for an exponential time.
var ErrSearchTreeTooLarge = errors.New("geo: the database's search tree has far more networks than nodes")

// OpenMMDBVerified is OpenMMDB with the verification bounded (code review
// R8-1). maxminddb's Verify walks every network path with no visited set,
// so a crafted file whose nodes share children enumerates 2^depth paths
// and never returns. A pre-pass over the same walk, capped at a small
// multiple of the node count (a real tree has about one network per node,
// times a few for the IPv4 aliases) and checking ctx, refuses such a file
// in time before Verify runs.
func OpenMMDBVerified(ctx context.Context, path string) (*MMDB, error) {
	r, err := maxminddb.Open(path)
	if err != nil {
		return nil, err
	}
	if err := boundedTreeWalk(ctx, r); err != nil {
		r.Close()
		return nil, err
	}
	if err := r.Verify(); err != nil {
		r.Close()
		return nil, err
	}
	return newMMDB(r), nil
}

// CheckRecords decodes every network's record through the schema live
// lookups use (code review R11-1). maxminddb's Verify decodes records
// generically, but a lookup decodes into mmdbRecord, which skips unknown
// fields with a recursion one frame per value: a record carrying a huge
// ignored array passes Verify and then overflows a serving process's stack
// on the first lookup in its network. Run ONLY in the verifier child, under
// its memory cap and deadline: there such a record kills the child, which
// refuses the database. Aliased IPv4 networks are skipped (same records).
func (m *MMDB) CheckRecords(ctx context.Context) error {
	if m == nil || m.r == nil {
		return errors.New("geo: no database")
	}
	limit := 4*m.nodes + 1024
	it := m.r.Networks(maxminddb.SkipAliasedNetworks)
	var n uint64
	for it.Next() {
		var rec mmdbRecord
		if _, err := it.Network(&rec); err != nil {
			return err
		}
		n++
		if n > limit {
			return ErrSearchTreeTooLarge
		}
		if n&0xffff == 0 {
			if err := ctx.Err(); err != nil {
				return err
			}
		}
	}
	return it.Err()
}

func boundedTreeWalk(ctx context.Context, r *maxminddb.Reader) error {
	limit := 4*uint64(r.Metadata.NodeCount) + 1024
	it := r.Networks()
	var n uint64
	for it.Next() {
		n++
		if n > limit {
			return ErrSearchTreeTooLarge
		}
		if n&0xffff == 0 {
			if err := ctx.Err(); err != nil {
				return err
			}
		}
	}
	return it.Err()
}

// OpenMMDBTrusted opens without Verify. Only for a file whose bytes were
// verified elsewhere and proven identical by digest (geoip spec §5.3):
// Verify decodes the whole data section, which costs seconds and tens of
// MiB on every edge for no new information.
func OpenMMDBTrusted(path string) (*MMDB, error) {
	r, err := maxminddb.Open(path)
	if err != nil {
		return nil, err
	}
	return newMMDB(r), nil
}

func newMMDB(r *maxminddb.Reader) *MMDB {
	md := r.Metadata
	return &MMDB{
		r: r, lookup: r.Lookup,
		source: sanitizeSource(md.DatabaseType),
		build:  time.Unix(int64(md.BuildEpoch), 0).UTC(),
		nodes:  uint64(md.NodeCount),
	}
}

// sanitizeSource reduces a metadata database_type to [A-Za-z0-9._-],
// at most 64 bytes: it is echoed into a request header and a log line.
func sanitizeSource(s string) string {
	var b strings.Builder
	for _, c := range s {
		if b.Len() >= 64 {
			break
		}
		if c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || c == '.' || c == '_' || c == '-' {
			b.WriteRune(c)
		}
	}
	if b.Len() == 0 {
		// Never empty (codex R7-2): a loaded reader's source must survive
		// the hop header's grammar, or a local hit becomes "off" downstream.
		return "mmdb"
	}
	return b.String()
}

// Source is the database's type ("DBIP-City-Lite", "GeoLite2-City"), sanitised.
func (m *MMDB) Source() string {
	if m == nil {
		return ""
	}
	return m.source
}

// BuildEpoch is when the database was built (UTC); zero for a nil MMDB.
func (m *MMDB) BuildEpoch() time.Time {
	if m == nil {
		return time.Time{}
	}
	return m.build
}

// NodeCount is the search tree's size, the shrink check's yardstick.
func (m *MMDB) NodeCount() uint64 {
	if m == nil {
		return 0
	}
	return m.nodes
}

type mmdbRecord struct {
	Location struct {
		Latitude  *float64 `maxminddb:"latitude"`
		Longitude *float64 `maxminddb:"longitude"`
	} `maxminddb:"location"`
	Country struct {
		ISOCode string `maxminddb:"iso_code"`
	} `maxminddb:"country"`
	RegisteredCountry struct {
		ISOCode string `maxminddb:"iso_code"`
	} `maxminddb:"registered_country"`
	Subdivisions []struct {
		ISOCode string `maxminddb:"iso_code"`
	} `maxminddb:"subdivisions"`
}

var (
	lookupErrMu   sync.Mutex
	lookupErrLast time.Time
)

// decode is the one place a record is read and judged (codex R1-9): the
// Locatable filter, IPv4 normalisation, the fixed-category error log, and
// the rule that a malformed location makes the WHOLE record a miss.
func (m *MMDB) decode(ip net.IP) (mmdbRecord, bool) {
	var rec mmdbRecord
	if m == nil || m.lookup == nil || !Locatable(ip) {
		return rec, false
	}
	if v4 := ip.To4(); v4 != nil {
		ip = v4
	}
	if err := m.lookup(ip, &rec); err != nil {
		lookupErrMu.Lock()
		if time.Since(lookupErrLast) > time.Minute {
			lookupErrLast = time.Now()
			log.Print("geo: a GeoIP lookup failed and was treated as a miss")
		}
		lookupErrMu.Unlock()
		return mmdbRecord{}, false
	}
	if rec.Location.Latitude != nil && rec.Location.Longitude != nil &&
		!onTheGlobe(Point{*rec.Location.Latitude, *rec.Location.Longitude}) {
		// A structurally valid database can still carry a malformed
		// location (final review). A record wrong in one field is not
		// trusted on the next.
		return mmdbRecord{}, false
	}
	return rec, true
}

// Locate returns the record's location (City databases), else its
// country's point, else its registered country's. Behaviour is exactly
// what it was before Lookup existed (codex R1-9).
func (m *MMDB) Locate(ip net.IP) (Point, bool) {
	rec, ok := m.decode(ip)
	if !ok {
		return Point{}, false
	}
	if rec.Location.Latitude != nil && rec.Location.Longitude != nil {
		return Point{*rec.Location.Latitude, *rec.Location.Longitude}, true
	}
	if p, ok := CountryPoint(rec.Country.ISOCode); ok {
		return p, true
	}
	return CountryPoint(rec.RegisteredCountry.ISOCode)
}

// Result is what Lookup found. Country is ISO 3166-1 alpha-2 (the
// record's country, else its registered country); Subdivision is
// "CC-XXX" only when the database carries an ISO code for it (GeoLite2
// does, DB-IP Lite does not, FACTS.md); HasPoint is true only for a real
// location record, never a country centroid.
type Result struct {
	Country     string
	Subdivision string
	Point       Point
	HasPoint    bool
}

// Lookup answers with whatever the record holds. ok is true when it found
// a country, a point, or both (geoip spec §6).
func (m *MMDB) Lookup(ip net.IP) (Result, bool) {
	rec, ok := m.decode(ip)
	if !ok {
		return Result{}, false
	}
	var r Result
	if isISO2(rec.Country.ISOCode) {
		r.Country = rec.Country.ISOCode
		if len(rec.Subdivisions) > 0 && isSubdivisionCode(rec.Subdivisions[0].ISOCode) {
			r.Subdivision = r.Country + "-" + rec.Subdivisions[0].ISOCode
		}
	} else if isISO2(rec.RegisteredCountry.ISOCode) {
		// The registered country says where the block is registered, not
		// which subdivision of it the visitor is in: no subdivision.
		r.Country = rec.RegisteredCountry.ISOCode
	}
	if rec.Location.Latitude != nil && rec.Location.Longitude != nil {
		r.Point, r.HasPoint = Point{*rec.Location.Latitude, *rec.Location.Longitude}, true
	}
	return r, r.Country != "" || r.HasPoint
}

func isISO2(s string) bool {
	return len(s) == 2 && s[0] >= 'A' && s[0] <= 'Z' && s[1] >= 'A' && s[1] <= 'Z'
}

func isSubdivisionCode(s string) bool {
	if len(s) < 1 || len(s) > 3 {
		return false
	}
	for i := 0; i < len(s); i++ {
		if !(s[i] >= 'A' && s[i] <= 'Z' || s[i] >= '0' && s[i] <= '9') {
			return false
		}
	}
	return true
}

// Close releases the mapping.
func (m *MMDB) Close() error {
	if m == nil {
		return nil
	}
	if m.closeHook != nil {
		m.closeHook()
	}
	if m.r == nil {
		return nil
	}
	return m.r.Close()
}

// onTheGlobe is false for nonfinite coordinates, a latitude outside
// [-90, 90] or a longitude outside [-180, 180]. Both bounds are inclusive:
// the poles and the antimeridian are places.
func onTheGlobe(p Point) bool {
	finite := func(f float64) bool { return !math.IsNaN(f) && !math.IsInf(f, 0) }
	return finite(p.Lat) && finite(p.Lon) && p.Lat >= -90 && p.Lat <= 90 && p.Lon >= -180 && p.Lon <= 180
}

var cgnat = &net.IPNet{IP: net.IPv4(100, 64, 0, 0), Mask: net.CIDRMask(10, 32)}

// Locatable is false for addresses that say nothing about where a person
// is: loopback, private, link-local, CGNAT and unspecified.
func Locatable(ip net.IP) bool {
	if ip == nil || ip.IsLoopback() || ip.IsPrivate() || ip.IsUnspecified() ||
		ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || cgnat.Contains(ip) {
		return false
	}
	return true
}