// Copyright 2019 CUE Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package net defines net-related types.
package net

import (
	"errors"
	"fmt"
	"math/big"
	"net/netip"

	"cuelang.org/go/cue"
)

// IP address lengths (bytes).
const (
	IPv4len = 4
	IPv6len = 16
)

func netGetIP(ip cue.Value) (netip.Addr, error) {
	switch ip.Kind() {
	case cue.StringKind:
		s, err := ip.String()
		if err != nil {
			return netip.Addr{}, err
		}
		goip, err := netip.ParseAddr(s)
		if err != nil {
			return netip.Addr{}, err
		}
		return goip, nil

	case cue.BytesKind:
		b, err := ip.Bytes()
		if err != nil {
			return netip.Addr{}, err
		}
		goip, err := netip.ParseAddr(string(b))
		if err != nil {
			return netip.Addr{}, err
		}
		return goip, nil

	case cue.ListKind:
		iter, err := ip.List()
		if err != nil {
			return netip.Addr{}, err
		}
		var bytes []byte
		for iter.Next() {
			v, err := iter.Value().Int64()
			if err != nil {
				return netip.Addr{}, err
			}
			if v < 0 || 255 < v {
				return netip.Addr{}, fmt.Errorf("invalid IP byte value %d", v)
			}
			bytes = append(bytes, byte(v))
		}
		goip, ok := netip.AddrFromSlice(bytes)
		if !ok {
			return netip.Addr{}, fmt.Errorf("invalid IP length %d", len(bytes))
		}
		return goip, nil

	default:
		return netip.Addr{}, fmt.Errorf("invalid type %v for IP address", ip.Kind())
	}
}

func netGetIPCIDR(ip cue.Value) (gonet *netip.Prefix, err error) {
	switch ip.Kind() {
	case cue.StringKind:
		s, err := ip.String()
		if err != nil {
			return nil, err
		}
		cidr, err := netip.ParsePrefix(s)
		if err != nil {
			return nil, err
		}
		return &cidr, nil

	case cue.BytesKind:
		b, err := ip.Bytes()
		if err != nil {
			return nil, err
		}
		cidr, err := netip.ParsePrefix(string(b))
		if err != nil {
			return nil, err
		}
		return &cidr, nil

	default:
		return nil, fmt.Errorf("invalid type %v for IP CIDR", ip.Kind())
	}
}

// ParseIP parses s as an IP address, returning the result.
// The string s can be in dotted decimal ("192.0.2.1")
// or IPv6 ("2001:db8::68") form.
// If s is not a valid textual representation of an IP address,
// ParseIP returns an error.
func ParseIP(s string) ([]uint, error) {
	goip, err := netip.ParseAddr(s)
	if err != nil {
		return nil, fmt.Errorf("invalid IP address %q", s)
	}
	return netToList(goip.AsSlice()), nil
}

func netToList(ip []byte) []uint {
	a := make([]uint, len(ip))
	for i, p := range ip {
		a[i] = uint(p)
	}
	return a
}

// IPv4 reports whether ip is a valid IPv4 address.
//
// The address may be a string or list of bytes.
func IPv4(ip cue.Value) (bool, error) {
	// TODO: convert to native CUE.
	goip, err := netGetIP(ip)
	if err != nil {
		return false, err
	}
	return goip.Is4(), nil
}

// IPv6 reports whether ip is a valid IPv6 address.
//
// The address may be a string or list of bytes.
func IPv6(ip cue.Value) (bool, error) {
	goip, err := netGetIP(ip)
	if err != nil {
		return false, err
	}
	return goip.Is6(), nil
}

// IP reports whether ip is a valid IPv4 or IPv6 address.
//
// The address may be a string or list of bytes.
func IP(ip cue.Value) (bool, error) {
	// TODO: convert to native CUE.
	goip, err := netGetIP(ip)
	if err != nil {
		return false, err
	}
	return goip.IsValid(), nil
}

// IPCIDR reports whether ip is a valid IPv4 or IPv6 address with CIDR subnet notation.
//
// The address may be a string or list of bytes.
func IPCIDR(ip cue.Value) (bool, error) {
	_, err := netGetIPCIDR(ip)
	return err == nil, err
}

// LoopbackIP reports whether ip is a loopback address.
func LoopbackIP(ip cue.Value) (bool, error) {
	goip, err := netGetIP(ip)
	if err != nil {
		return false, err
	}
	return goip.IsLoopback(), nil
}

// MulticastIP reports whether ip is a multicast address.
func MulticastIP(ip cue.Value) (bool, error) {
	goip, err := netGetIP(ip)
	if err != nil {
		return false, err
	}
	return goip.IsMulticast(), nil
}

// InterfaceLocalMulticastIP reports whether ip is an interface-local multicast
// address.
func InterfaceLocalMulticastIP(ip cue.Value) (bool, error) {
	goip, err := netGetIP(ip)
	if err != nil {
		return false, err
	}
	return goip.IsInterfaceLocalMulticast(), nil
}

// LinkLocalMulticastIP reports whether ip is a link-local multicast address.
func LinkLocalMulticastIP(ip cue.Value) (bool, error) {
	goip, err := netGetIP(ip)
	if err != nil {
		return false, err
	}
	return goip.IsLinkLocalMulticast(), nil
}

// LinkLocalUnicastIP reports whether ip is a link-local unicast address.
func LinkLocalUnicastIP(ip cue.Value) (bool, error) {
	goip, err := netGetIP(ip)
	if err != nil {
		return false, err
	}
	return goip.IsLinkLocalUnicast(), nil
}

// GlobalUnicastIP reports whether ip is a global unicast address.
//
// The identification of global unicast addresses uses address type
// identification as defined in RFC 1122, RFC 4632 and RFC 4291 with the
// exception of IPv4 directed broadcast addresses. It returns true even if ip is
// in IPv4 private address space or local IPv6 unicast address space.
func GlobalUnicastIP(ip cue.Value) (bool, error) {
	goip, err := netGetIP(ip)
	if err != nil {
		return false, err
	}
	return goip.IsGlobalUnicast(), nil
}

// UnspecifiedIP reports whether ip is an unspecified address, either the IPv4
// address "0.0.0.0" or the IPv6 address "::".
func UnspecifiedIP(ip cue.Value) (bool, error) {
	goip, err := netGetIP(ip)
	if err != nil {
		return false, err
	}
	return goip.IsUnspecified(), nil
}

// ToIP4 converts a given IP address, which may be a string or a list, to its
// 4-byte representation.
func ToIP4(ip cue.Value) ([]uint, error) {
	ipdata, err := netGetIP(ip)
	if err != nil {
		return nil, err
	}
	if !ipdata.Is4() {
		return nil, fmt.Errorf("cannot convert %q to IPv4", ipdata)
	}
	as4 := ipdata.As4()
	return netToList(as4[:]), nil
}

// ToIP16 converts a given IP address, which may be a string or a list, to its
// 16-byte representation.
func ToIP16(ip cue.Value) ([]uint, error) {
	ipdata, err := netGetIP(ip)
	if err != nil {
		return nil, err
	}
	as16 := ipdata.As16()
	return netToList(as16[:]), nil
}

// IPString returns the string form of the IP address ip. It returns one of 4 forms:
//
// - "<nil>", if ip has length 0
// - dotted decimal ("192.0.2.1"), if ip is an IPv4 or IP4-mapped IPv6 address
// - IPv6 ("2001:db8::1"), if ip is a valid IPv6 address
// - the hexadecimal form of ip, without punctuation, if no other cases apply
func IPString(ip cue.Value) (string, error) {
	ipdata, err := netGetIP(ip)
	if err != nil {
		return "", err
	}
	return ipdata.String(), nil
}

func netIPAdd(addr netip.Addr, offset *big.Int) (netip.Addr, error) {
	i := big.NewInt(0).SetBytes(addr.AsSlice())
	i = i.Add(i, offset)

	if i.Sign() < 0 {
		return netip.Addr{}, errors.New("IP address arithmetic resulted in out-of-range address (underflow)")
	}

	b := i.Bytes()
	size := addr.BitLen() / 8

	if len(b) > size {
		return netip.Addr{}, errors.New("IP address arithmetic resulted in out-of-range address (overflow)")
	}

	if len(b) < size {
		b = append(make([]byte, size-len(b), size), b...)
	}
	addr, _ = netip.AddrFromSlice(b)
	return addr, nil
}

// AddIP adds a numerical offset to a given IP address.
// The address can be provided as a string, byte array, or CIDR subnet notation.
// It returns the resulting IP address or CIDR subnet notation as a string.
func AddIP(ip cue.Value, offset *big.Int) (string, error) {
	prefix, err := netGetIPCIDR(ip)
	if err == nil {
		addr, err := netIPAdd(prefix.Addr(), offset)
		if err != nil {
			return "", err
		}
		return netip.PrefixFrom(addr, prefix.Bits()).String(), nil
	}
	ipdata, err := netGetIP(ip)
	if err != nil {
		return "", err
	}
	addr, err := netIPAdd(ipdata, offset)
	if err != nil {
		return "", err
	}
	return addr.String(), nil
}

// AddIPCIDR adds a numerical offset to a given CIDR subnet
// string, returning a CIDR string.
func AddIPCIDR(ip cue.Value, offset *big.Int) (string, error) {
	prefix, err := netGetIPCIDR(ip)
	if err != nil {
		return "", err
	}
	shifted := big.NewInt(0).Lsh(offset, (uint)(prefix.Addr().BitLen()-prefix.Bits()))
	addr, err := netIPAdd(prefix.Addr(), shifted)
	if err != nil {
		return "", err
	}
	return netip.PrefixFrom(addr, prefix.Bits()).String(), nil
}

// ParsedCIDR holds the parsed components of a CIDR notation string.
type ParsedCIDR struct {
	PrefixMask string `json:"prefix_mask"`
	PrefixLen  int    `json:"prefix_len"`
	PrefixAddr string `json:"prefix_addr"`
	// BroadcastAddr is only set for IPv4 CIDRs.
	BroadcastAddr string `json:"broadcast_addr,omitempty"`
}

// ParseCIDR parses a CIDR notation string and returns its components:
// prefix_mask (e.g. "255.255.255.0"), prefix_len (e.g. 24),
// prefix_addr (e.g. "10.20.30.0"), and broadcast_addr (e.g. "10.20.30.255").
// broadcast_addr is only set for IPv4 CIDRs.
func ParseCIDR(s string) (*ParsedCIDR, error) {
	prefix, err := netip.ParsePrefix(s)
	if err != nil {
		return nil, err
	}

	bits := prefix.Bits()
	addr := prefix.Addr()
	maskBytes := make([]byte, addr.BitLen()/8)
	for i := range bits / 8 {
		maskBytes[i] = 0xFF
	}
	if rem := bits % 8; rem > 0 {
		maskBytes[bits/8] = ^byte(0xFF >> rem)
	}
	netmask, _ := netip.AddrFromSlice(maskBytes)

	networkAddr := prefix.Masked().Addr()

	result := &ParsedCIDR{
		PrefixMask: netmask.String(),
		PrefixLen:  bits,
		PrefixAddr: networkAddr.String(),
	}

	if addr.Is4() {
		broadcastBytes := networkAddr.AsSlice()
		for i := range broadcastBytes {
			broadcastBytes[i] |= ^maskBytes[i]
		}
		broadcastAddr, _ := netip.AddrFromSlice(broadcastBytes)
		result.BroadcastAddr = broadcastAddr.String()
	}

	return result, nil
}

// InCIDR reports whether an IP address is contained a CIDR subnet string.
func InCIDR(ip, cidr cue.Value) (bool, error) {
	ipAddr, err := netGetIP(ip)
	if err != nil {
		return false, err
	}

	prefix, err := netGetIPCIDR(cidr)
	if err != nil {
		return false, err
	}

	return prefix.Contains(ipAddr), nil
}

// CompareIP compares two IP addresses and returns an integer:
// -1 if ip1 sorts before ip2, 0 if they are equal, and +1 if ip1 sorts after ip2.
// IPv4 addresses sort before IPv6 addresses.
//
// The addresses may be strings or lists of bytes.
func CompareIP(ip1, ip2 cue.Value) (int, error) {
	addr1, err := netGetIP(ip1)
	if err != nil {
		return 0, err
	}
	addr2, err := netGetIP(ip2)
	if err != nil {
		return 0, err
	}
	return addr1.Compare(addr2), nil
}
