package main

import (
	"encoding/binary"
	"net/netip"
	"sort"
)

type addressRange struct {
	start uint32
	end   uint32
}

func rangeFromPrefix(value string) (addressRange, bool) {
	prefix, err := netip.ParsePrefix(value)
	if err != nil || !prefix.Addr().Is4() {
		return addressRange{}, false
	}
	address := prefix.Masked().Addr().As4()
	start := binary.BigEndian.Uint32(address[:])
	hostBits := 32 - prefix.Bits()
	end := start
	if hostBits == 32 {
		end = ^uint32(0)
	} else if hostBits > 0 {
		end |= (uint32(1) << hostBits) - 1
	}
	return addressRange{start: start, end: end}, true
}

func mergeRanges(ranges []addressRange) []addressRange {
	if len(ranges) == 0 {
		return nil
	}
	sort.Slice(ranges, func(i, j int) bool {
		if ranges[i].start == ranges[j].start {
			return ranges[i].end < ranges[j].end
		}
		return ranges[i].start < ranges[j].start
	})
	merged := []addressRange{ranges[0]}
	for _, candidate := range ranges[1:] {
		current := &merged[len(merged)-1]
		adjacent := current.end != ^uint32(0) && candidate.start == current.end+1
		if candidate.start <= current.end || adjacent {
			if candidate.end > current.end {
				current.end = candidate.end
			}
			continue
		}
		merged = append(merged, candidate)
	}
	return merged
}

func intersectRanges(left, right []addressRange) []addressRange {
	left = mergeRanges(left)
	right = mergeRanges(right)
	var result []addressRange
	for leftIndex, rightIndex := 0, 0; leftIndex < len(left) && rightIndex < len(right); {
		start := max(left[leftIndex].start, right[rightIndex].start)
		end := min(left[leftIndex].end, right[rightIndex].end)
		if start <= end {
			result = append(result, addressRange{start: start, end: end})
		}
		if left[leftIndex].end < right[rightIndex].end {
			leftIndex++
		} else {
			rightIndex++
		}
	}
	return mergeRanges(result)
}

func subtractRanges(base, denied []addressRange) []addressRange {
	base = mergeRanges(base)
	denied = mergeRanges(denied)
	var result []addressRange
	denyIndex := 0
	for _, candidate := range base {
		cursor := uint64(candidate.start)
		end := uint64(candidate.end)
		for denyIndex < len(denied) && uint64(denied[denyIndex].end) < cursor {
			denyIndex++
		}
		for index := denyIndex; index < len(denied) && uint64(denied[index].start) <= end; index++ {
			if uint64(denied[index].start) > cursor {
				result = append(result, addressRange{start: uint32(cursor), end: denied[index].start - 1})
			}
			if uint64(denied[index].end)+1 > cursor {
				cursor = uint64(denied[index].end) + 1
			}
			if cursor > end {
				break
			}
		}
		if cursor <= end {
			result = append(result, addressRange{start: uint32(cursor), end: uint32(end)})
		}
	}
	return mergeRanges(result)
}

func addressCount(ranges []addressRange) uint64 {
	var result uint64
	for _, current := range mergeRanges(ranges) {
		result += uint64(current.end) - uint64(current.start) + 1
	}
	return result
}