/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

package netaddr

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"strconv"
)

type NetAddr struct {
	Address []byte
	Length  byte
}

func applyMask(address, mask []byte) []byte {
	var ip net.IP = net.IP(address)

	return ip.Mask(mask)
}

func (na *NetAddr) AddressSize() int {
	return len(na.Address)
}

func (na *NetAddr) Mask() []byte {
	onesToMask := [9]byte{0, 128, 192, 224, 240, 248, 252, 254, 255}
	min := func(a, b byte) byte {
		if a < b {
			return a
		}
		return b
	}
	ones := na.Length
	mask := make([]byte, 0, na.AddressSize())
	for i := 0; i < na.AddressSize(); i++ {
		olen := min(8, ones)
		ones -= olen
		mask = append(mask, onesToMask[int(olen)])
	}
	return mask
}

func (na *NetAddr) FirstAddress() *NetAddr {
	bitSize := na.AddressSize() * 8
	length := na.Length
	mask := net.CIDRMask(int(length), bitSize)
	addressableSlice := na.Address
	first0 := applyMask(addressableSlice, mask)
	first := make([]byte, na.AddressSize())

	copy(first[:], first0)

	return &NetAddr{Address: first, Length: length}
}

func (na *NetAddr) LastAddress() *NetAddr {
	bitSize := na.AddressSize() * 8
	length := na.Length
	mask := net.CIDRMask(int(length), bitSize)
	allOnes := net.CIDRMask(bitSize, bitSize)
	invMask := byteXor(mask, allOnes)
	lastAddr := byteOr(na.FirstAddress().Address, invMask)

	return &NetAddr{Address: lastAddr, Length: length}
}

func (na *NetAddr) String() string {
	addr := net.IP(na.Address).String()
	length := fmt.Sprintf("%d", na.Length)
	return addr + "/" + length
}

func (na *NetAddr) IsRFC1918() bool {
	return IsSubset(na, &NetAddr{
		Address: []byte{10, 0, 0, 0},
		Length:  byte(8),
	}) || IsSubset(na, &NetAddr{
		Address: []byte{172, 16, 0, 0},
		Length:  byte(12),
	}) || IsSubset(na, &NetAddr{
		Address: []byte{192, 168, 0, 0},
		Length:  byte(16),
	})
}

func byteOp(u, v []byte, f func(byte, byte) byte) []byte {
	w := make([]byte, len(u), len(u))

	for i := 0; i < len(u); i++ {
		w[i] = f(u[i], v[i])
	}
	return w
}

func byteOr(u, v []byte) []byte {
	return byteOp(u, v, func(ui, vi byte) byte {
		return ui | vi
	})
}

func byteXor(u, v []byte) []byte {
	return byteOp(u, v, func(ui, vi byte) byte {
		return ui ^ vi
	})
}

func SetLength(na *NetAddr, length byte) *NetAddr {
	min := func(a, b byte) byte {
		if a < b {
			return a
		}
		return b
	}
	return &NetAddr{
		Address: na.Address,
		Length:  min(length, byte(na.AddressSize()*8)),
	}
}

/* Test against skewed total order on byte-size prefixes */
func octetLessThan(o1, m1, o2, m2 byte) bool {
	f1 := int(o1&m1) - int((o1^byte(255))&m1)
	f2 := int(o2&m2) - int((o2^byte(255))&m2)
	return f1 < f2 && !(o1&m1 == o2&m1 && m1 <= m2)
}

func LessThan(na1 *NetAddr, na2 *NetAddr) (less bool) {
	if na1.AddressSize() != na2.AddressSize() {
		return
	}
	mask1 := na1.Mask()
	mask2 := na2.Mask()
	for i := 0; i < na1.AddressSize(); i++ {
		if octetLessThan(
			na1.Address[i], mask1[i],
			na2.Address[i], mask2[i],
		) {
			less = true
			break
		}
		if na1.Address[i] == na2.Address[i] &&
			mask1[i] == mask2[i] {
			continue
		}
		break
	}
	return
}

func IsEqual(na1, na2 *NetAddr) bool {
	return na1.Length == na2.Length &&
		bytes.Equal(na1.Address, na2.Address)
}

func Contains(na1, na2 *NetAddr) bool {
	tmp := SetLength(na2, na1.Length)

	return na1.Length <= na2.Length &&
		IsEqual(na1.FirstAddress(), tmp.FirstAddress())
}

func IsSubset(na1, na2 *NetAddr) bool {
	tmp := SetLength(na1, na2.Length)

	return na2.Length <= na1.Length &&
		IsEqual(na2.FirstAddress(), tmp.FirstAddress())
}

func Overlaps(na1, na2 *NetAddr) bool {
	return Contains(na1, na2) || Contains(na2, na1)
}

func parseTruncatedIPv4(s string) (na *NetAddr, err error) {
	na = &NetAddr{Address: make([]byte, 4)}
	r := bytes.NewBufferString(s)
	buf := new(bytes.Buffer)
	i := 0
	onLen := false
	for i < 4 {
		var b byte
		b, err = r.ReadByte()
		switch {
		case err == io.EOF:
			na.Length = byte(8*i - 8)
			if onLen {
				var o int
				o, err = strconv.Atoi(buf.String())
				if err != nil {
					return
				}
				if o < 0 || 32 < o {
					err = fmt.Errorf("length out of range: %v", o)
					return
				}
				na.Length = byte(o)
				return
			}
			err = nil
			return
		case err != nil:
			return
		case '0' <= b && b <= '9':
			buf.WriteByte(b)
			continue
		case b == '.' || b == '/':
			var o int
			o, err = strconv.Atoi(buf.String())
			if err != nil {
				return
			}
			if o < 0 || 255 < o {
				err = fmt.Errorf("octet out of range: %v", o)
				return
			}
			na.Address[i] = byte(o)
			i++
			buf = new(bytes.Buffer)
			if b == '/' {
				onLen = true
			}
		}
	}
	err = fmt.Errorf("input is longer than expected")
	return
}

func IP(s string) (na *NetAddr, err error) {
	var ip net.IP
	var ipNet *net.IPNet

	ip, ipNet, err = net.ParseCIDR(s)
	switch err.(type) {
	case nil:
	case *net.ParseError:
		ip = net.ParseIP(s)
		if ip == nil {
			na, err = parseTruncatedIPv4(s)
			if err != nil {
				err = fmt.Errorf("invalid ip address %v: %v", s, err)
			}
			return
		}
		if test := ip.To4(); test != nil {
			ipNet = &net.IPNet{
				IP:   ip,
				Mask: net.IPMask([]byte{255, 255, 255, 255}),
			}
		} else if test := ip.To16(); test != nil {
			ipNet = &net.IPNet{
				IP: ip,
				Mask: net.IPMask(
					[]byte{
						0xff, 0xff, 0xff, 0xff,
						0xff, 0xff, 0xff, 0xff,
						0xff, 0xff, 0xff, 0xff,
						0xff, 0xff, 0xff, 0xff,
					},
				),
			}
		}
		err = nil
	default:
		return
	}
	var length, _ int = net.IPMask.Size(ipNet.Mask)
	var address []byte = ip.To4()

	if address == nil {
		address = ip.To16()
	}
	if address == nil {
		panic(fmt.Sprintf("unexpected failure parsing '%s'\n", s))
	}
	na = &NetAddr{Address: address, Length: byte(length)}
	return
}