package ipmetagen

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"go/format"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/klauspost/compress/zstd"
)

const (
	cityAssetFile = "data_city.hk.zst"
	asnAssetFile  = "data_asn.hk.zst"
)

type binHeader struct {
	databaseType      uint8
	databaseColumn    uint8
	ipv4DatabaseCount uint32
	ipv4DatabaseAddr  uint32
	ipv6DatabaseCount uint32
	ipv6DatabaseAddr  uint32
}

type cityMeta struct {
	countryCode string
	region      string
	city        string
}

type asnMeta struct {
	asn      int
	provider string
}

type packedRangeBuilder[T comparable] struct {
	metaIDs     map[T]uint32
	metaRecords []T
	ipv4Starts  []byte
	ipv4Meta    []byte
	ipv6Starts  []byte
	ipv6Meta    []byte
	lastIPv4ID  uint32
	lastIPv6ID  uint32
	hasIPv4     bool
	hasIPv6     bool
}

// WritePackedCityBINAssets writes compact city lookup assets from an
// IP2Location LITE DB3 BIN database.
func WritePackedCityBINAssets(dir string, data []byte) error {
	header, err := readBINHeader(data)
	if err != nil {
		return err
	}
	if header.databaseType != 3 {
		return fmt.Errorf("expected DB3 BIN database type 3, got %d", header.databaseType)
	}
	builder := newPackedRangeBuilder[cityMeta]()
	if err := walkBINRows(data, header, func(family AddressFamily, start []byte, row []byte) error {
		meta := cityMeta{
			countryCode: cleanCountryCode(readBINString(data, readBINPointer(row, 0))),
			region:      cleanField(readBINString(data, readBINPointer(row, 4))),
			city:        cleanField(readBINString(data, readBINPointer(row, 8))),
		}
		if err := builder.add(family, start, meta); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return err
	}
	return writePackedLookupAsset(dir, cityAssetFile, "HKCY", builder.ipv4Starts, builder.ipv4Meta, builder.ipv6Starts, builder.ipv6Meta, renderCityMeta(builder.metaRecords))
}

// WritePackedASNBINAssets writes compact ASN lookup assets from an IP2Location
// LITE ASN BIN database.
func WritePackedASNBINAssets(dir string, data []byte) error {
	header, err := readBINHeader(data)
	if err != nil {
		return err
	}
	if header.databaseType != 26 {
		return fmt.Errorf("expected ASN BIN database type 26, got %d", header.databaseType)
	}
	builder := newPackedRangeBuilder[asnMeta]()
	if err := walkBINRows(data, header, func(family AddressFamily, start []byte, row []byte) error {
		asn := 0
		if rawASN := cleanField(readBINString(data, readBINPointer(row, 88))); rawASN != "" {
			parsed, err := strconv.Atoi(rawASN)
			if err != nil {
				return fmt.Errorf("parse ASN %q: %w", rawASN, err)
			}
			asn = parsed
		}
		meta := asnMeta{
			asn:      asn,
			provider: cleanField(readBINString(data, readBINPointer(row, 92))),
		}
		if err := builder.add(family, start, meta); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return err
	}
	return writePackedLookupAsset(dir, asnAssetFile, "HKAS", builder.ipv4Starts, builder.ipv4Meta, builder.ipv6Starts, builder.ipv6Meta, renderASNMeta(builder.metaRecords))
}

// WritePackedCityNetworkEmbedSource writes the generated Go source that embeds
// compact city and ASN lookup assets.
func WritePackedCityNetworkEmbedSource(dir string, packageName string) error {
	if strings.TrimSpace(packageName) == "" {
		packageName = "ipmeta"
	}
	var b strings.Builder
	fmt.Fprintf(&b, "// Code generated by go run ./cmd/ipmeta-generate; DO NOT EDIT.\n")
	fmt.Fprintf(&b, "//\n")
	fmt.Fprintf(&b, "// HitKeep uses the IP2Location LITE database for IP geolocation.\n")
	fmt.Fprintf(&b, "// https://www.ip2location.com\n")
	fmt.Fprintf(&b, "package %s\n\n", packageName)
	fmt.Fprintf(&b, "import _ \"embed\"\n\n")
	fmt.Fprintf(&b, "//go:embed %s\n", cityAssetFile)
	fmt.Fprintf(&b, "var embeddedCityZSTDData []byte\n\n")
	fmt.Fprintf(&b, "//go:embed %s\n", asnAssetFile)
	fmt.Fprintf(&b, "var embeddedASNZSTDData []byte\n")

	source, err := format.Source([]byte(b.String()))
	if err != nil {
		return fmt.Errorf("format city/network embed source: %w", err)
	}
	if err := writePublicGeneratedFile(filepath.Join(dir, "data_city_network_lite.go"), source); err != nil {
		return fmt.Errorf("write city/network embed source: %w", err)
	}
	return nil
}

func readBINHeader(data []byte) (binHeader, error) {
	if len(data) < 64 {
		return binHeader{}, errors.New("IP2Location BIN is shorter than the 64-byte header")
	}
	return binHeader{
		databaseType:      data[0],
		databaseColumn:    data[1],
		ipv4DatabaseCount: binary.LittleEndian.Uint32(data[5:9]),
		ipv4DatabaseAddr:  binary.LittleEndian.Uint32(data[9:13]),
		ipv6DatabaseCount: binary.LittleEndian.Uint32(data[13:17]),
		ipv6DatabaseAddr:  binary.LittleEndian.Uint32(data[17:21]),
	}, nil
}

func walkBINRows(data []byte, header binHeader, visit func(AddressFamily, []byte, []byte) error) error {
	ipv4ColumnSize := int(header.databaseColumn) * 4
	ipv6ColumnSize := 16 + ((int(header.databaseColumn) - 1) * 4)
	if err := walkBINFamily(data, IPv4, header.ipv4DatabaseAddr, header.ipv4DatabaseCount, ipv4ColumnSize, 4, visit); err != nil {
		return err
	}
	return walkBINFamily(data, IPv6, header.ipv6DatabaseAddr, header.ipv6DatabaseCount, ipv6ColumnSize, 16, visit)
}

func walkBINFamily(data []byte, family AddressFamily, addr uint32, count uint32, columnSize int, firstColumnSize int, visit func(AddressFamily, []byte, []byte) error) error {
	if addr == 0 || count == 0 {
		return nil
	}
	base := int(addr) - 1
	if base < 0 {
		return fmt.Errorf("invalid %s database address %d", family, addr)
	}
	for index := 0; index < int(count); index++ {
		offset := base + (index * columnSize)
		end := offset + columnSize
		if offset < 0 || end > len(data) {
			return fmt.Errorf("%s row %d exceeds BIN size", family, index)
		}
		fullRow := data[offset:end]
		start := normalizeBINStart(family, fullRow[:firstColumnSize])
		row := fullRow[firstColumnSize:]
		if err := visit(family, start, row); err != nil {
			return fmt.Errorf("%s row %d: %w", family, index, err)
		}
	}
	return nil
}

func normalizeBINStart(family AddressFamily, raw []byte) []byte {
	out := make([]byte, len(raw))
	switch family {
	case IPv4:
		value := binary.LittleEndian.Uint32(raw)
		binary.BigEndian.PutUint32(out, value)
	case IPv6:
		copy(out, raw)
		reverseBytes(out)
	default:
		copy(out, raw)
	}
	return out
}

func reverseBytes(data []byte) {
	for i, j := 0, len(data)-1; i < j; i, j = i+1, j-1 {
		data[i], data[j] = data[j], data[i]
	}
}

func newPackedRangeBuilder[T comparable]() *packedRangeBuilder[T] {
	return &packedRangeBuilder[T]{metaIDs: map[T]uint32{}}
}

func (b *packedRangeBuilder[T]) add(family AddressFamily, start []byte, meta T) error {
	id, ok := b.metaIDs[meta]
	if !ok {
		var err error
		id, err = checkedUint32(len(b.metaRecords), "metadata records")
		if err != nil {
			return err
		}
		b.metaIDs[meta] = id
		b.metaRecords = append(b.metaRecords, meta)
	}
	idBytes := make([]byte, 4)
	binary.BigEndian.PutUint32(idBytes, id)
	switch family {
	case IPv4:
		if b.hasIPv4 && b.lastIPv4ID == id {
			return nil
		}
		b.ipv4Starts = append(b.ipv4Starts, start...)
		b.ipv4Meta = append(b.ipv4Meta, idBytes...)
		b.lastIPv4ID = id
		b.hasIPv4 = true
	case IPv6:
		if b.hasIPv6 && b.lastIPv6ID == id {
			return nil
		}
		b.ipv6Starts = append(b.ipv6Starts, start...)
		b.ipv6Meta = append(b.ipv6Meta, idBytes...)
		b.lastIPv6ID = id
		b.hasIPv6 = true
	}
	return nil
}

func readBINPointer(row []byte, offset int) uint32 {
	if offset < 0 || offset+4 > len(row) {
		return 0
	}
	return binary.LittleEndian.Uint32(row[offset : offset+4])
}

func readBINString(data []byte, pointer uint32) string {
	if pointer == 0 || pointer > math.MaxInt32 {
		return ""
	}
	offset := int(pointer)
	if offset < 0 || offset >= len(data) {
		return ""
	}
	length := int(data[offset])
	start := offset + 1
	end := start + length
	if end > len(data) {
		return ""
	}
	return string(data[start:end])
}

func writePackedLookupAsset(dir string, name string, magic string, ipv4Starts []byte, ipv4Meta []byte, ipv6Starts []byte, ipv6Meta []byte, meta []byte) error {
	if len(magic) != 4 {
		return fmt.Errorf("asset magic %q must be 4 bytes", magic)
	}
	var raw bytes.Buffer
	raw.WriteString(magic)
	ipv4Count, err := checkedRangeCount(ipv4Starts, 4, "IPv4 ranges")
	if err != nil {
		return err
	}
	ipv6Count, err := checkedRangeCount(ipv6Starts, 16, "IPv6 ranges")
	if err != nil {
		return err
	}
	metaSize, err := checkedUint32(len(meta), "metadata payload")
	if err != nil {
		return err
	}
	writeUint32(&raw, ipv4Count)
	writeUint32(&raw, ipv6Count)
	writeUint32(&raw, metaSize)
	raw.Write(ipv4Starts)
	raw.Write(ipv4Meta)
	raw.Write(ipv6Starts)
	raw.Write(ipv6Meta)
	raw.Write(meta)

	compressed, err := zstdBytes(raw.Bytes())
	if err != nil {
		return fmt.Errorf("compress asset %s: %w", name, err)
	}
	if err := writePublicGeneratedFile(filepath.Join(dir, name), compressed); err != nil {
		return fmt.Errorf("write asset %s: %w", name, err)
	}
	return nil
}

func checkedRangeCount(data []byte, stride int, label string) (uint32, error) {
	if stride <= 0 {
		return 0, fmt.Errorf("%s stride must be positive", label)
	}
	if len(data)%stride != 0 {
		return 0, fmt.Errorf("%s data length %d is not divisible by stride %d", label, len(data), stride)
	}
	return checkedUint32(len(data)/stride, label)
}

func checkedUint32(value int, label string) (uint32, error) {
	if value < 0 || value > math.MaxUint32 {
		return 0, fmt.Errorf("%s count %d exceeds uint32", label, value)
	}
	//nolint:gosec // value is bounded above by math.MaxUint32 before conversion.
	return uint32(value), nil
}

func writePublicGeneratedFile(path string, data []byte) error {
	if err := os.WriteFile(path, data, 0o600); err != nil {
		return err
	}
	return os.Chmod(path, 0o644)
}

func writeUint32(w *bytes.Buffer, value uint32) {
	var raw [4]byte
	binary.BigEndian.PutUint32(raw[:], value)
	w.Write(raw[:])
}

func zstdBytes(data []byte) ([]byte, error) {
	writer, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.SpeedBestCompression))
	if err != nil {
		return nil, err
	}
	defer writer.Close()
	return writer.EncodeAll(data, nil), nil
}

func renderCityMeta(records []cityMeta) []byte {
	lines := make([]string, 0, len(records))
	for _, record := range records {
		lines = append(lines, strings.Join([]string{
			escapeMetadataField(record.countryCode),
			escapeMetadataField(record.region),
			escapeMetadataField(record.city),
		}, "\t"))
	}
	return []byte(strings.Join(lines, "\n"))
}

func renderASNMeta(records []asnMeta) []byte {
	lines := make([]string, 0, len(records))
	for _, record := range records {
		lines = append(lines, strconv.Itoa(record.asn)+"\t"+escapeMetadataField(record.provider))
	}
	return []byte(strings.Join(lines, "\n"))
}

func escapeMetadataField(value string) string {
	value = strings.ReplaceAll(value, "\t", " ")
	value = strings.ReplaceAll(value, "\n", " ")
	return value
}

func CompactAssetStats(dir string) (map[string]int64, error) {
	stats := map[string]int64{}
	for _, pattern := range []string{"data_city.*", "data_asn.*"} {
		matches, err := filepath.Glob(filepath.Join(dir, pattern))
		if err != nil {
			return nil, err
		}
		sort.Strings(matches)
		for _, path := range matches {
			info, err := os.Stat(path)
			if err != nil {
				return nil, err
			}
			stats[filepath.Base(path)] = info.Size()
		}
	}
	return stats, nil
}
