package pipeline

import (
	"fmt"
	"math/big"
	"math/rand"
	"net/netip"
	"sort"
	"strings"
	"testing"

	"evobgp/internal/store"
)

// The functions below are a verbatim copy of the pre-2.1 prune+merge loop
// (from git HEAD internal/pipeline/refresh.go). They exist only as an
// equivalence oracle for collapsePrefixGroup.

func legacyAggregateCIDRGroup(rows []store.PrefixRow, mergeFn func(map[string]store.PrefixRow) bool) []store.PrefixRow {
	if len(rows) <= 1 {
		return rows
	}
	set := make(map[string]store.PrefixRow, len(rows))
	for _, row := range rows {
		set[row.Prefix] = row
	}
	legacyPruneCoveredPrefixes(set)
	for {
		if !mergeFn(set) {
			break
		}
		legacyPruneCoveredPrefixes(set)
	}
	out := make([]store.PrefixRow, 0, len(set))
	for _, row := range set {
		out = append(out, row)
	}
	sortPrefixRows(out)
	return out
}

func legacyPruneCoveredPrefixes(set map[string]store.PrefixRow) {
	type item struct {
		key  string
		pfx  netip.Prefix
		bits int
	}
	items := make([]item, 0, len(set))
	for k := range set {
		p, err := netip.ParsePrefix(k)
		if err != nil {
			continue
		}
		items = append(items, item{key: k, pfx: p, bits: p.Bits()})
	}
	sort.Slice(items, func(i, j int) bool {
		if items[i].bits != items[j].bits {
			return items[i].bits < items[j].bits
		}
		return items[i].key < items[j].key
	})
	for i := 0; i < len(items); i++ {
		for j := i + 1; j < len(items); j++ {
			if items[j].bits <= items[i].bits {
				continue
			}
			if items[i].pfx.Contains(items[j].pfx.Addr()) {
				delete(set, items[j].key)
			}
		}
	}
}

func legacyMergeSiblingPrefixesIPv4(set map[string]store.PrefixRow) bool {
	merged := false
	seen := make(map[string]struct{}, len(set))
	for key, row := range set {
		if _, done := seen[key]; done {
			continue
		}
		pfx, err := netip.ParsePrefix(key)
		if err != nil || !pfx.Addr().Is4() {
			continue
		}
		bits := pfx.Bits()
		if bits <= 8 {
			continue
		}
		netNum := ipv4PrefixNetworkLegacy(pfx)
		blockSize := uint32(1) << (32 - bits)
		siblingNet := netNum ^ blockSize
		siblingPfx := netip.PrefixFrom(u32ToIPv4(siblingNet), bits).Masked().String()
		if _, ok := set[siblingPfx]; !ok {
			continue
		}
		parentBits := bits - 1
		parentBlock := uint32(1) << (32 - parentBits)
		parentNet := netNum & ^(parentBlock - 1)
		parentPfx := netip.PrefixFrom(u32ToIPv4(parentNet), parentBits).Masked().String()
		delete(set, key)
		delete(set, siblingPfx)
		parentRow := row
		parentRow.Prefix = parentPfx
		set[parentPfx] = parentRow
		seen[key] = struct{}{}
		seen[siblingPfx] = struct{}{}
		merged = true
	}
	return merged
}

func ipv4PrefixNetworkLegacy(p netip.Prefix) uint32 {
	a := p.Masked().Addr().As4()
	return uint32(a[0])<<24 | uint32(a[1])<<16 | uint32(a[2])<<8 | uint32(a[3])
}

func legacyMergeSiblingPrefixesIPv6(set map[string]store.PrefixRow) bool {
	merged := false
	seen := make(map[string]struct{}, len(set))
	for key, row := range set {
		if _, done := seen[key]; done {
			continue
		}
		pfx, err := netip.ParsePrefix(key)
		if err != nil || !pfx.Addr().Is6() {
			continue
		}
		bits := pfx.Bits()
		if bits <= 16 {
			continue
		}
		netNum := ipv6PrefixNetworkLegacy(pfx)
		blockSize := new(big.Int).Lsh(big.NewInt(1), uint(128-bits))
		siblingNet := new(big.Int).Xor(netNum, blockSize)
		siblingPfx := ipv6PrefixFromBigIntLegacy(siblingNet, bits).String()
		if _, ok := set[siblingPfx]; !ok {
			continue
		}
		parentBits := bits - 1
		parentBlock := new(big.Int).Lsh(big.NewInt(1), uint(128-parentBits))
		mask := new(big.Int).Sub(parentBlock, big.NewInt(1))
		mask.Not(mask)
		parentNet := new(big.Int).And(netNum, mask)
		parentPfx := ipv6PrefixFromBigIntLegacy(parentNet, parentBits).String()
		delete(set, key)
		delete(set, siblingPfx)
		parentRow := row
		parentRow.Prefix = parentPfx
		set[parentPfx] = parentRow
		seen[key] = struct{}{}
		seen[siblingPfx] = struct{}{}
		merged = true
	}
	return merged
}

func ipv6PrefixNetworkLegacy(p netip.Prefix) *big.Int {
	a := p.Masked().Addr().As16()
	n := new(big.Int)
	n.SetBytes(a[:])
	return n
}

func ipv6PrefixFromBigIntLegacy(n *big.Int, bits int) netip.Prefix {
	b := n.Bytes()
	var a [16]byte
	copy(a[16-len(b):], b)
	return netip.PrefixFrom(netip.AddrFrom16(a), bits).Masked()
}

func normalizeForCompare(rows []store.PrefixRow) []string {
	type line struct{ p, c, s string }
	lines := make([]line, 0, len(rows))
	for _, r := range rows {
		p := strings.TrimSpace(r.Prefix)
		if p != "" {
			if pfx, err := netip.ParsePrefix(p); err == nil {
				p = pfx.Masked().String()
			}
		}
		lines = append(lines, line{p, prefixRowCommunity(r), r.Source})
	}
	sort.Slice(lines, func(i, j int) bool {
		if lines[i].p != lines[j].p {
			return lines[i].p < lines[j].p
		}
		if lines[i].c != lines[j].c {
			return lines[i].c < lines[j].c
		}
		return lines[i].s < lines[j].s
	})
	out := make([]string, 0, len(lines))
	for _, l := range lines {
		out = append(out, fmt.Sprintf("%s|%s|%s", l.p, l.c, l.s))
	}
	return out
}

func TestCollapsePrefixGroup_EquivalenceWithLegacy(t *testing.T) {
	r := rand.New(rand.NewSource(7))
	for iter := 0; iter < 200; iter++ {
		n := 1 + r.Intn(60)
		v4rows := make([]store.PrefixRow, 0, n)
		for i := 0; i < n; i++ {
			addr := netip.AddrFrom4([4]byte{203, byte(r.Intn(4)), byte(r.Intn(256)), byte(r.Intn(256))})
			bits := 16 + r.Intn(9)
			pfx := netip.PrefixFrom(addr, bits).Masked()
			v4rows = append(v4rows, store.PrefixRow{Prefix: pfx.String(), Source: "ip_range"})
		}
		legacy := legacyAggregateCIDRGroup(append([]store.PrefixRow(nil), v4rows...), legacyMergeSiblingPrefixesIPv4)
		got := collapsePrefixGroup(append([]store.PrefixRow(nil), v4rows...), true)
		if fmt.Sprint(normalizeForCompare(legacy)) != fmt.Sprint(normalizeForCompare(got)) {
			t.Fatalf("iter %d mismatch:\nlegacy=%v\ngot   =%v", iter, normalizeForCompare(legacy), normalizeForCompare(got))
		}
	}
}

func TestCollapsePrefixGroup_EquivalenceWithLegacyIPv6(t *testing.T) {
	r := rand.New(rand.NewSource(11))
	for iter := 0; iter < 200; iter++ {
		n := 1 + r.Intn(60)
		rows := make([]store.PrefixRow, 0, n)
		for i := 0; i < n; i++ {
			var a [16]byte
			a[0], a[1] = 0x20, 0x01
			a[2], a[3] = 0x0d, 0xb8
			a[4] = byte(r.Intn(2))
			a[5] = byte(r.Intn(256))
			a[6] = byte(r.Intn(256))
			addr := netip.AddrFrom16(a)
			bits := 32 + r.Intn(17)
			pfx := netip.PrefixFrom(addr, bits).Masked()
			rows = append(rows, store.PrefixRow{Prefix: pfx.String(), Source: "ip_range"})
		}
		legacy := legacyAggregateCIDRGroup(append([]store.PrefixRow(nil), rows...), legacyMergeSiblingPrefixesIPv6)
		got := collapsePrefixGroup(append([]store.PrefixRow(nil), rows...), false)
		if fmt.Sprint(normalizeForCompare(legacy)) != fmt.Sprint(normalizeForCompare(got)) {
			t.Fatalf("iter %d mismatch:\nlegacy=%v\ngot   =%v", iter, normalizeForCompare(legacy), normalizeForCompare(got))
		}
	}
}

func TestCollapsePrefixGroup_CoversAndSiblingChain(t *testing.T) {
	rows := []store.PrefixRow{
		{Prefix: "10.0.0.0/16", Source: "ip_range"},
		{Prefix: "10.0.0.0/24", Source: "ip_range"},
		{Prefix: "10.0.1.0/24", Source: "ip_range"},
		{Prefix: "10.0.2.0/24", Source: "ip_range"},
		{Prefix: "10.0.3.0/24", Source: "ip_range"},
		{Prefix: "10.1.0.0/24", Source: "ip_range"},
	}
	got := collapsePrefixGroup(rows, true)
	if len(got) != 2 {
		t.Fatalf("want 2 rows (/16 + /24), got %d: %+v", len(got), got)
	}
	set := map[string]bool{}
	for _, r := range got {
		set[r.Prefix] = true
	}
	if !set["10.0.0.0/16"] || !set["10.1.0.0/24"] {
		t.Fatalf("unexpected rows: %+v", got)
	}
}

func TestCollapsePrefixGroup_RespectsMinBitsFloor(t *testing.T) {
	rows := []store.PrefixRow{
		{Prefix: "10.0.0.0/8", Source: "ip_range"},
		{Prefix: "11.0.0.0/8", Source: "ip_range"},
	}
	got := collapsePrefixGroup(rows, true)
	if len(got) != 2 {
		t.Fatalf("floor must prevent /8+/8 -> /7, got %+v", got)
	}

	rows9 := []store.PrefixRow{
		{Prefix: "10.0.0.0/9", Source: "ip_range"},
		{Prefix: "10.128.0.0/9", Source: "ip_range"},
	}
	got9 := collapsePrefixGroup(rows9, true)
	if len(got9) != 1 || got9[0].Prefix != "10.0.0.0/8" {
		t.Fatalf("expected /9+/9 -> /8, got %+v", got9)
	}
}

func TestCollapsePrefixGroup_PrunesCovered(t *testing.T) {
	rows := []store.PrefixRow{
		{Prefix: "10.0.0.0/16", Source: "ip_range"},
		{Prefix: "10.0.0.0/24", Source: "ip_range"},
	}
	got := collapsePrefixGroup(rows, true)
	if len(got) != 1 || got[0].Prefix != "10.0.0.0/16" {
		t.Fatalf("want covered prune to /16, got %+v", got)
	}
}