// Package lookuptest writes lookup databases, IPinfo Lite files in their
// .mmdb form, for the tests of the packages that read them.
package lookuptest

import (
	"bytes"
	"net"
	"os"
	"testing"

	"github.com/maxmind/mmdbwriter"
	"github.com/maxmind/mmdbwriter/mmdbtype"
)

// fileMode is the mode of the files written: read and written by their
// owner alone.
const fileMode = 0o600

// Network is what a lookup database holds about a netblock, of the fields
// smallwebwaf reads: its AS number, such as AS64496, the AS's name, and
// its country, such as DE.
type Network struct {
	ASN     string
	ASName  string
	Country string
}

// Write writes a lookup database at path that places each netblock in
// networks, such as 203.0.113.0/24, as its Network says, and no other
// address.
func Write(tb testing.TB, path string, networks map[string]Network) {
	tb.Helper()

	records := make(map[string]mmdbtype.Map, len(networks))
	for netblock, network := range networks {
		records[netblock] = mmdbtype.Map{
			"asn":          mmdbtype.String(network.ASN),
			"as_name":      mmdbtype.String(network.ASName),
			"country_code": mmdbtype.String(network.Country),
		}
	}

	WriteRecords(tb, path, records)
}

// WriteRecords writes a lookup database at path that holds each record in
// records for its netblock, and nothing for any other address.
func WriteRecords(tb testing.TB, path string, records map[string]mmdbtype.Map) {
	tb.Helper()

	tree, err := mmdbwriter.New(mmdbwriter.Options{
		DatabaseType: "ipinfo_lite",
		// The tests' clients are in the netblocks kept for documentation.
		IncludeReservedNetworks: true,
	})
	if err != nil {
		tb.Fatalf("new lookup database: %v", err)
	}

	for netblock, record := range records {
		_, network, err := net.ParseCIDR(netblock)
		if err != nil {
			tb.Fatalf("netblock %q: %v", netblock, err)
		}

		err = tree.Insert(network, record)
		if err != nil {
			tb.Fatalf("insert %s: %v", netblock, err)
		}
	}

	var database bytes.Buffer

	_, err = tree.WriteTo(&database)
	if err != nil {
		tb.Fatalf("write the lookup database: %v", err)
	}

	err = os.WriteFile(path, database.Bytes(), fileMode)
	if err != nil {
		tb.Fatalf("write %s: %v", path, err)
	}
}