package netsplit

import (
	`bytes`
	"encoding/json"
	"encoding/xml"
	"fmt"
	"math"
	"math/big"
	"net"
	"net/netip"
	"strings"

	"github.com/goccy/go-yaml"
	"go4.org/netipx"
)

/*
AddrExpand expands a netip.Addr's string format.
Like netip.Addr.StringExpanded() but for IPv4 too.
*/
func AddrExpand(ip netip.Addr) (s string) {

	var sb *strings.Builder

	if ip.IsUnspecified() || !ip.IsValid() {
		return
	}

	if ip.Is6() {
		s = ip.StringExpanded()
	} else {
		// IPv4 we have to do by hand.
		sb = new(strings.Builder)
		for idx, b := range ip.AsSlice() {
			sb.WriteString(fmt.Sprintf("%03d", b))
			if idx != net.IPv4len-1 {
				sb.WriteString(".")
			}
		}
		s = sb.String()
	}

	return
}

/*
AddrCompress returns the shortest possible CIDR representation as a string from a netip.Prefix.
Note that IPv6 netip.Prefix.String() already does this automatically, as IPv6 has special condensing rules.
*/
func AddrCompress(pfx *netip.Prefix) (s string) {

	var sl []string
	var lastNonzero int

	if pfx == nil || !pfx.IsValid() || !pfx.Addr().IsValid() {
		return
	}

	if pfx.Addr().Is6() {
		s = pfx.String()
		return
	}

	sl = strings.Split(pfx.Addr().String(), ".")

	for idx, oct := range sl {
		if oct != "0" {
			lastNonzero = idx
		}
	}

	s = fmt.Sprintf("%s/%d", strings.Join(sl[:lastNonzero+1], "."), pfx.Bits())

	return
}

/*
AddrFmt provides a string representation for an IP (as a netip.Addr).

`f` is the string formatter to use (without the %). For IPv4, you generally want `d`,
for IPv6, you generally want `x`.

`sep` indicates a character to insert every `every` bytes of the mask.
For IPv4, you probably want `.`,
for IPv6 there isn't really a standard representation; CIDR notation is preferred.
Thus for IPv6 you probably want to set sep as blank and/or set `every` to 0.

`segSep` indicates a character sequence to use for segmenting the string.
Specify as an empty string and/or set `everySeg` to 0 to disable.

`every` indicates how many bytes should pass before sep is inserted.
For IPv4, this should be 1.
For IPv6, there isn't really a standard indication but it's recommended to do 2.
Set as 0 or `sep` to an empty string to do no separation characters.

`everySeg` indicates how many *seperations* should pass before segSep is inserted.
Set as 0 or `segSep` to an empty string to do no string segmentation.
*/
func AddrFmt(ip netip.Addr, f, sep, segSep string, every, everySeg uint) (s string) {

	var numSegs int
	var doSep = every > 0
	var fs = "%" + f
	var sb = new(strings.Builder)

	if ip.IsUnspecified() || !ip.IsValid() {
		return
	}
	for idx, b := range ip.AsSlice() {
		if doSep && idx > 0 {
			if idx%int(every) == 0 {
				sb.WriteString(sep)
				numSegs++
			}
			if everySeg > 0 {
				if numSegs >= int(everySeg) {
					sb.WriteString(segSep)
					numSegs = 0
				}
			}
		}
		fmt.Fprintf(sb, fs, b)
	}

	s = strings.TrimSpace(sb.String())

	return
}

/*
AddrInvert returns an inverted form of netip.Addr as another netip.Addr.

Note that it doesn't really make sense to use this for IPv6.
*/
func AddrInvert(ip netip.Addr) (inverted netip.Addr) {

	var b []byte

	if !ip.IsValid() {
		return
	}

	b = make([]byte, len([]byte(ip.AsSlice())))

	for idx, i := range []byte(ip.AsSlice()) {
		b[idx] = ^i
	}

	inverted, _ = netip.AddrFromSlice(b)

	return
}

/*
CheckReserved checks nets for any reserved prefixes; either directly/explicitly,
included *within* a reserved prefix (revRecursive), or *including* a reserved prefix (recursive).
excludePrivate indicates if LAN networks should be considered as "reserved" or not.
If a network is found via revRecursive/recursive, the matching prefix - not the specified one - will be in reservations.

Any found will be returned in reservations.

If no reserved networks are found, reservations will be nil.

Note that prefix-specific broadcasts (e.g. x.255.255.255/8, x.x.x.255/24, ::/64, x:ffff:ffff:ffff:ffff/64, etc.)
will *not* be considered as "reserved" as they are considered normal addresses expected for functionality.
This primarily focuses on prefixes/subnets for this reason.
Additionally, all of nets will be aligned to their proper boundary range/CIDR/subnet.
*/
func CheckReserved(nets []*netip.Prefix, revRecursive, recursive, excludePrivate bool) (reservations map[netip.Prefix]*IANAAddrNetResRecord, err error) {

	var ok bool
	var res *IANAAddrNetResRecord
	var reserved map[netip.Prefix]*IANAAddrNetResRecord

	if nets == nil || len(nets) == 0 {
		return
	}

	if _, _, reserved, err = RetrieveReserved(); err != nil {
		return
	}

	for _, n := range nets {
		if n == nil {
			continue
		}
		if n.Addr().IsPrivate() && excludePrivate {
			continue
		}
		*n = n.Masked()
		if res, ok = reserved[*n]; ok {
			if reservations == nil {
				reservations = make(map[netip.Prefix]*IANAAddrNetResRecord)
			}
			reservations[*n] = res
		}
		if !revRecursive && !recursive {
			continue
		}
		for p, r := range reserved {
			// This... *should* be safe? I don't think any reservations overlap.
			// Anyways, revRecursive works because n.Addr() returns the network address, which should be the canonical boundary.
			// recursive works for the same reason, just the other end.
			// Math!
			if revRecursive && p.Contains(n.Addr()) {
				if reservations == nil {
					reservations = make(map[netip.Prefix]*IANAAddrNetResRecord)
				}
				reservations[p] = r
			}
			if recursive && n.Bits() < p.Bits() && n.Contains(p.Addr()) {
				if reservations == nil {
					reservations = make(map[netip.Prefix]*IANAAddrNetResRecord)
				}
				reservations[p] = r
			}
		}
	}

	return
}

// Contain takes the results of a NetSplitter and returns a StructuredResults. The reservations are only checked against nets.
func Contain(origPfx *netip.Prefix, nets []*netip.Prefix, remaining *netipx.IPSet, splitter NetSplitter) (s *StructuredResults, err error) {

	var rem []netip.Prefix
	// var reserved map[netip.Prefix]*IANAAddrNetResRecord
	var sr = StructuredResults{
		Original: origPfx,
	}

	if origPfx == nil {
		return
	}

	if origPfx.Addr() != origPfx.Masked().Addr() {
		sr.Canonical = new(netip.Prefix)
		*sr.Canonical = origPfx.Masked()
		sr.HostAddr = new(netip.Addr)
		*sr.HostAddr = origPfx.Addr()
	}

	if splitter != nil {
		sr.Splitter = new(SplitOpts)
		switch t := splitter.(type) {
		case *CIDRSplitter:
			sr.Splitter.CIDR = t
		case *HostSplitter:
			sr.Splitter.Host = t
		case *SubnetSplitter:
			sr.Splitter.Subnet = t
		case *VLSMSplitter:
			sr.Splitter.VLSM = t
		default:
			err = ErrBadSplitter
			return
		}
	}

	if nets != nil {
		sr.Allocated = make([]*ContainedResult, len(nets))
		for idx, n := range nets {
			sr.Allocated[idx] = &ContainedResult{
				Network: n,
			}
		}
	}

	if remaining != nil {
		rem = remaining.Prefixes()
		sr.Unallocated = make([]*ContainedResult, len(rem))
		for idx, i := range rem {
			sr.Unallocated[idx] = &ContainedResult{
				Network: new(netip.Prefix),
			}
			*sr.Unallocated[idx].Network = i
		}
	}

	/*
		if nets != nil {
			if reserved, err = CheckReserved(nets, true, true, false); err != nil {
				return
			}
			if reserved != nil && len(reserved) > 0 {
				s.Reservations = make([]*IANAAddrNetResRecord, len(reserved))
				idx := 0
				for _, r := range reserved {
					s.Reservations[idx] = r
					idx++
				}
			}
		}
	*/

	s = &sr

	return
}

/*
LastSubnetNet returns the last subnet of prefix length pfxLen from network pfx.

lastSubnet will be nil if pfx is nil or invalid.
*/
func LastSubnetNet(pfx *net.IPNet, pfxLen uint8) (lastSubnet *net.IPNet, err error) {

	var nPfx netip.Prefix
	var ok bool

	if pfx == nil {
		return
	}
	if nPfx, ok = netipx.FromStdIPNet(pfx); !ok {
		return
	}
	if !nPfx.IsValid() {
		return
	}

	if nPfx, err = LastSubnetPfx(nPfx, pfxLen); err != nil {
		return
	}
	if !nPfx.IsValid() {
		return
	}

	lastSubnet = netipx.PrefixIPNet(nPfx)

	return
}

// LastSubnetPfx is exactly like LastSubnetNet but using netip.Prefix instead.
func LastSubnetPfx(pfx netip.Prefix, pfxLen uint8) (lastSubnet netip.Prefix, err error) {

	var ok bool
	var pfxBytes []byte
	var bitOffset uint
	var offset *big.Int
	var ipInt *big.Int
	var ipBytes []byte
	var byteLen int
	var lastNet *net.IPNet
	var maxBitLen uint8 = maxBitsv4

	if !pfx.IsValid() {
		return
	}

	pfx = pfx.Masked()

	if pfx.Addr().Is6() || pfxLen > maxBitsv4 {
		maxBitLen = maxBitsv6
	}
	if pfxLen > maxBitLen {
		err = ErrBadPrefixLen
		return
	}
	if pfxLen < uint8(pfx.Bits()) {
		err = ErrBigPrefix
		return
	}

	if pfxBytes, err = pfx.Addr().MarshalBinary(); err != nil {
		return
	}

	byteLen = pfx.Addr().BitLen() / 8
	bitOffset = uint(pfxLen - uint8(pfx.Bits()))

	offset = new(big.Int).Lsh(big.NewInt(1), bitOffset)
	offset.Sub(offset, big.NewInt(1))

	// Cast the prefix (as represented as bytes) into a *big.Int for some number magic.
	ipInt = new(big.Int).SetBytes(pfxBytes)
	// Shift to the first "address" in the prefix/mask it.
	offset.Lsh(offset, uint(pfx.Addr().BitLen()-int(pfxLen)))
	// Now add the offset to the base network.
	ipInt.Add(ipInt, offset)

	// If the base address starts at the "0 address" (e.g. 0.0.0.0 or :: etc....),
	// this can cause some strange behavior when casting the *big.Int to bytes.
	// So it gets left-null-padded to the appropriate length for the inet family.
	ipBytes = ipInt.Bytes()
	if len(ipBytes) < byteLen {
		ipBytes = append(
			bytes.Repeat([]byte{0x00}, byteLen-len(ipBytes)),
			ipBytes...,
		)
	}

	// Create an explicit net.IPNet...
	lastNet = &net.IPNet{
		IP:   net.IP(ipBytes),
		Mask: net.CIDRMask(int(pfxLen), pfx.Addr().BitLen()),
	}
	// And then make it a netip.Prefix.
	if lastSubnet, ok = netipx.FromStdIPNet(lastNet); !ok {
		return
	}
	if !lastSubnet.IsValid() {
		return
	}

	return
}

/*
MaskExpand expands a net.IPMask's string format.
Like AddrExpand but for netmasks.
*/
func MaskExpand(mask net.IPMask, isIpv6 bool) (s string) {

	var sb *strings.Builder

	// IPv6 is always expanded in string format, but not split out.
	if isIpv6 {
		s = MaskFmt(mask, "02x", ":", "", 2, 0)
		return
	}

	sb = new(strings.Builder)
	for idx, b := range mask {
		sb.WriteString(fmt.Sprintf("%03d", b))
		if idx != net.IPv4len-1 {
			sb.WriteString(".")
		}
	}
	s = sb.String()

	return
}

/*
MaskFmt provides a string representation for a netmask (as a net.IPMask).

Its parameters hold the same significance as in AddrFmt.
*/
func MaskFmt(mask net.IPMask, f, sep, segSep string, every, everySeg uint) (s string) {

	var numSegs int
	var doSep = every > 0
	var fs = "%" + f
	var sb = new(strings.Builder)

	if mask == nil || len(mask) == 0 {
		return
	}
	for idx, b := range mask {
		if doSep && idx > 0 {
			if idx%int(every) == 0 {
				sb.WriteString(sep)
				numSegs++
			}
			if everySeg > 0 {
				if numSegs >= int(everySeg) {
					sb.WriteString(segSep)
					numSegs = 0
				}
			}
		}
		fmt.Fprintf(sb, fs, b)
	}

	s = strings.TrimSpace(sb.String())

	return
}

/*
MaskInvert returns an inverted form of net.IPMask as another net.IPMask.

Note that it doesn't really make sense to use this for IPv6.
*/
func MaskInvert(mask net.IPMask) (inverted net.IPMask) {

	var b []byte

	b = make([]byte, len([]byte(mask)))

	for idx, i := range []byte(mask) {
		b[idx] = ^i
	}

	inverted = net.IPMask(b)

	return
}

/*
	NumAddrsIn returns the number of addresses in a given prefix length
	and inet family.

	If isIpv6 is false, it is assumed to be IPv4 (...duh).

	inclNet and inclBcast have the same meanings as in NumAddrsNet and NumAddrsPfx.

	Note that for the single-host prefix (/32 for IPv4, /128 for IPv6), numAddrs will *always* be 1.
	For point-to-point prefix (IPv4 /31, IPv6 /127), numAddrs will *ALWAYS* be 2.
*/
func NumAddrsIn(pfxLen uint8, isIpv6, inclNet, inclBcast bool) (numAddrs *big.Int, err error) {

	var numBits uint
	var numRemoved int64
	var maxBitLen uint8 = maxBitsv4

	if isIpv6 || pfxLen > maxBitsv4 {
		maxBitLen = maxBitsv6
	}
	if pfxLen > maxBitLen {
		err = ErrBadPrefixLen
		return
	}
	if pfxLen == maxBitLen {
		numAddrs = big.NewInt(1)
		return
	}
	if (pfxLen + 1) == maxBitLen {
		numAddrs = big.NewInt(2)
		return
	}

	numBits = uint(maxBitLen - pfxLen)

	numAddrs = new(big.Int).Lsh(big.NewInt(1), numBits)

	if !inclNet {
		numRemoved++
	}
	if !inclBcast {
		numRemoved++
	}
	if numRemoved > 0 {
		_ = numAddrs.Sub(numAddrs, big.NewInt(numRemoved))
	}

	return
}

/*
NumAddrsNet returns the number of IP addresses in a net.IPNet.

The network address is included in the count if inclNet is true, otherwise it is excluded.

The broadcast (or reserved broadcast, in the case of IPv6) address will be included in
the count if inclBcast is true, otherwise it is excluded.

numAddrs will be nil if pfx is nil or invalid.
*/
func NumAddrsNet(pfx *net.IPNet, inclNet, inclBcast bool) (numAddrs *big.Int) {

	var nPfx netip.Prefix
	var ok bool

	if pfx == nil {
		return
	}
	if nPfx, ok = netipx.FromStdIPNet(pfx); !ok {
		return
	}
	if !nPfx.IsValid() {
		return
	}

	numAddrs = NumAddrsPfx(nPfx, inclNet, inclBcast)

	return
}

// NumAddrsPfx is the exact same as NumAddrsNet but for a net/netip.Prefix instead.
func NumAddrsPfx(pfx netip.Prefix, inclNet, inclBcast bool) (numAddrs *big.Int) {

	var err error

	if !pfx.IsValid() {
		return
	}

	// Since we're dealing with existing prefixes/networks, we should never get an error.
	if numAddrs, err = NumAddrsIn(uint8(pfx.Bits()), pfx.Addr().Is6(), inclNet, inclBcast); err != nil {
		// But *somehow* in case we do...
		panic(err)
	}

	return
}

/*
NumNets returns the number of times prefix size subnet fits into prefix size network.

It will error if network is larger than 128 or if subnet is smaller than network.

This is MUCH more performant than splitting out an actual network into explicit subnets,
and does not require an actual network.
*/
func NumNets(subnet, network uint8) (numNets uint, ipv6Only bool, err error) {

	var x float64

	// network cannot be higher than 128, as that's the maximum for IPv6.
	if network > maxBits {
		err = ErrBadPrefixLen
		return
	}
	if subnet < network {
		err = ErrBigPrefix
		return
	}
	ipv6Only = (network > maxBitsv4) || (subnet > maxBitsv4)

	x = float64(subnet - network)

	numNets = uint(math.Pow(2, x))

	return
}

// Parse parses b for JSON/XML/YAML and tries to return a StructuredResults from it.
func Parse(b []byte) (s *StructuredResults, err error) {

	if b == nil {
		return
	}

	if err = json.Unmarshal(b, &s); err != nil {
		if err = xml.Unmarshal(b, &s); err != nil {
			if err = yaml.Unmarshal(b, &s); err != nil {
				return
			} else {
				return
			}
		} else {
			return
		}
	}

	return
}

/*
ValidateSizes ensures that none of the prefix lengths in sizes exceeds the maximum possible in pfx.
No-ops with nil error if pfx is nil, sizes is nil, or sizes is empty.

err is also nil if validation succeeds.
If validation fails on a prefix length size, the error will be a SplitErr
with only Wrapped and RequestedPrefixLen fields populated *for the first failing size only*.
*/
func ValidateSizes(pfx *net.IPNet, sizes ...uint8) (err error) {

	var ok bool
	var addr netip.Addr
	var familyMax uint8

	if pfx == nil || sizes == nil || len(sizes) == 0 {
		return
	}
	if addr, ok = netipx.FromStdIP(pfx.IP); !ok {
		err = ErrBadPrefix
		return
	}
	if addr.Is4() {
		familyMax = 32
	} else {
		familyMax = 128
	}
	for _, size := range sizes {
		if size > familyMax {
			err = &SplitErr{
				Wrapped:            ErrBadPrefixLen,
				Nets:               nil,
				Remaining:          nil,
				LastSubnet:         nil,
				RequestedPrefixLen: size,
			}
			return
		}
	}

	return
}