package pipeline

import (
	"net/netip"
	"sort"

	"evobgp/internal/store"
)

// Prefix collapse replaces the former prune+merge fixed-point loop (O(n²) per pass,
// multiple passes) with a single sort + linear stack pass: O(n log n) overall.
//
// Classic trie-collapse without building a trie:
//  1. Sort prefixes by (address, mask len ascending).
//  2. Walk left to right keeping a stack of "open" prefixes:
//     - while the stack top contains the new prefix -> the top covers it, drop the new one (prune);
//     - else, while the stack top is a sibling of the new prefix (same mask, XOR of the
//     network bit equals the parent block) and merging is allowed for that mask length
//     -> pop the sibling, replace the new prefix with the parent and re-check;
//     - otherwise push the new prefix.
//
// A parent absorbs its sibling children only when both halves are present, which matches
// the previous merge-to-fixed-point semantics, including the guard that forbids merging
// above a floor mask (IPv4 /8, IPv6 /16).

type collapseItem struct {
	row store.PrefixRow
	pfx netip.Prefix
}

// collapsePrefixGroup collapses one (community, source) group of same-family prefixes.
// All rows must be masked; family and floor are enforced by minBits.
func collapsePrefixGroup(rows []store.PrefixRow, is4 bool) []store.PrefixRow {
	if len(rows) <= 1 {
		return rows
	}
	items := make([]collapseItem, 0, len(rows))
	seen := make(map[string]struct{}, len(rows))
	for _, row := range rows {
		if _, dup := seen[row.Prefix]; dup {
			continue
		}
		seen[row.Prefix] = struct{}{}
		pfx, err := netip.ParsePrefix(row.Prefix)
		if err != nil {
			continue
		}
		pfx = pfx.Masked()
		if is4 != pfx.Addr().Is4() {
			continue
		}
		items = append(items, collapseItem{row: row, pfx: pfx})
	}
	if len(items) <= 1 {
		out := make([]store.PrefixRow, 0, len(items))
		for _, it := range items {
			out = append(out, it.row)
		}
		return out
	}

	// Sort by (address, mask len): a parent always sorts before its children, and a
	// shorter sibling sorts before a longer one within the same parent block.
	sort.Slice(items, func(i, j int) bool {
		a, b := items[i].pfx, items[j].pfx
		if a.Addr() != b.Addr() {
			return lessAddr(a.Addr(), b.Addr())
		}
		return a.Bits() < b.Bits()
	})

	minBits := 8
	if !is4 {
		minBits = 16
	}

	stack := make([]collapseItem, 0, len(items))
	for _, it := range items {
		cur := it
		dropped := false
		for len(stack) > 0 {
			top := stack[len(stack)-1]
			if top.pfx.Contains(cur.pfx.Addr()) && top.pfx.Bits() <= cur.pfx.Bits() {
				// Covered by an existing prefix: drop (prune).
				dropped = true
				break
			}
			if cur.pfx.Bits() == top.pfx.Bits() && cur.pfx.Bits() > minBits && areSiblings(top.pfx, cur.pfx) {
				// Merge siblings into the parent (parent keeps the lower sibling's attributes),
				// then re-check the parent against the new stack top.
				stack = stack[:len(stack)-1]
				parentBits := cur.pfx.Bits() - 1
				parent := netip.PrefixFrom(maskAddr(cur.pfx.Addr(), parentBits, is4), parentBits).Masked()
				cur = collapseItem{row: top.row, pfx: parent}
				cur.row.Prefix = parent.String()
				continue
			}
			break
		}
		if !dropped {
			stack = append(stack, cur)
		}
	}

	out := make([]store.PrefixRow, 0, len(stack))
	for _, it := range stack {
		out = append(out, it.row)
	}
	sortPrefixRows(out)
	return out
}

// areSiblings reports whether two same-length prefixes combine into their common parent.
func areSiblings(a, b netip.Prefix) bool {
	if a.Bits() != b.Bits() || a.Bits() == 0 {
		return false
	}
	parentBits := a.Bits() - 1
	pa := maskAddr(a.Addr(), parentBits, a.Addr().Is4())
	pb := maskAddr(b.Addr(), parentBits, b.Addr().Is4())
	return pa == pb
}

// maskAddr clears the host bits below prefixLen (IPv4: 32-bit space; IPv6: 128-bit).
func maskAddr(a netip.Addr, prefixLen int, is4 bool) netip.Addr {
	if is4 {
		v := uint32FromIPv4(a)
		if prefixLen <= 0 {
			v = 0
		} else if prefixLen < 32 {
			v &= ^(uint32(1)<<(32-prefixLen) - 1)
		}
		return u32ToIPv4(v)
	}
	b := a.As16()
	hostBits := 128 - prefixLen
	fullBytes := hostBits / 8
	for i := 15; i > 15-fullBytes; i-- {
		b[i] = 0
	}
	if rem := hostBits % 8; rem > 0 {
		idx := 15 - fullBytes
		if idx >= 0 && idx < 16 {
			b[idx] &= byte(0xFF << rem)
		}
	}
	return netip.AddrFrom16(b)
}

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

func u32ToIPv4(v uint32) netip.Addr {
	return netip.AddrFrom4([4]byte{byte(v >> 24), byte(v >> 16), byte(v >> 8), byte(v)})
}

func lessAddr(a, b netip.Addr) bool {
	if a.Is4() != b.Is4() {
		return a.Is4()
	}
	if a.Is4() {
		return uint32FromIPv4(a) < uint32FromIPv4(b)
	}
	a16, b16 := a.As16(), b.As16()
	for i := 0; i < 16; i++ {
		if a16[i] != b16[i] {
			return a16[i] < b16[i]
		}
	}
	return false
}