package function

import (
	"errors"
	"net/netip"

	"github.com/brimdata/super"
	"github.com/brimdata/super/scode"
	"github.com/brimdata/super/sup"
)

type NetworkOf struct {
	sctx *super.Context
}

func (n *NetworkOf) Call(args []super.Value) super.Value {
	ipVal := args[0].Under()
	if ipVal.Type().ID() != super.IDIP {
		return n.sctx.WrapError("network_of: not an IP", ipVal)
	}
	ip := super.DecodeIP(ipVal.Bytes())
	var bits int
	if len(args) == 1 {
		switch {
		case !ip.Is4():
			return n.sctx.WrapError("network_of: not an IPv4 address", ipVal)
		case ip.As4()[0] < 0x80:
			bits = 8
		case ip.As4()[0] < 0xc0:
			bits = 16
		default:
			bits = 24
		}
	} else {
		// two args
		maskVal := args[1].Under()
		switch id := maskVal.Type().ID(); {
		case id == super.IDIP:
			mask := super.DecodeIP(maskVal.Bytes())
			if mask.BitLen() != ip.BitLen() {
				val := addressAndMask(n.sctx, ipVal, maskVal)
				return n.sctx.WrapError("network_of: address and mask have different lengths", val)
			}
			bits = super.LeadingOnes(mask.AsSlice())
			if netip.PrefixFrom(mask, bits).Masked().Addr() != mask {
				return n.sctx.WrapError("network_of: mask is non-contiguous", maskVal)
			}
		case super.IsInteger(id):
			if super.IsSigned(id) {
				bits = int(maskVal.Int())
			} else {
				bits = int(maskVal.Uint())
			}
			if bits > 128 || bits > 32 && ip.Is4() {
				val := addressAndMask(n.sctx, ipVal, maskVal)
				return n.sctx.WrapError("network_of: CIDR bit count out of range", val)
			}
		default:
			return n.sctx.WrapError("network_of: bad arg for CIDR mask", args[1])
		}
	}
	// Mask for canonical form.
	prefix := netip.PrefixFrom(ip, bits).Masked()
	return super.NewNet(prefix)
}

func addressAndMask(sctx *super.Context, address, mask super.Value) super.Value {
	val, err := sup.NewBSUPMarshalerWithContext(sctx).Marshal(struct {
		Address super.Value `super:"address"`
		Mask    super.Value `super:"mask"`
	}{address, mask})
	if err != nil {
		panic(err)
	}
	return val
}

type CIDRMatch struct {
	sctx *super.Context
}

var errMatch = errors.New("match")

func (c *CIDRMatch) Call(args []super.Value) super.Value {
	maskVal := args[0].Under()
	if id := maskVal.Type().ID(); id != super.IDNet && id != super.IDNull {
		val := c.sctx.WrapError("cidr_match: not a net", maskVal)
		if maskVal.IsNull() {
			val = super.NewValue(val.Type(), nil)
		}
		return val

	}
	if maskVal.IsNull() || args[1].IsNull() {
		return super.NewValue(super.TypeBool, nil)
	}
	prefix := super.DecodeNet(maskVal.Bytes())
	err := args[1].Walk(func(typ super.Type, body scode.Bytes) error {
		if typ.ID() == super.IDIP {
			if prefix.Contains(super.DecodeIP(body)) {
				return errMatch
			}
		}
		return nil
	})
	return super.NewBool(err == errMatch)
}
