package btcwire

import (
	"encoding/binary"
	"fmt"
	"strconv"
	"strings"
)

// NetAddress is a Bitcoin network address as carried by the version message:
// the services offered by the peer, its IP address, and its port.
//
// The IP field holds 16 bytes. IPv4 addresses are stored in their
// IPv4-mapped IPv6 form (::ffff:a.b.c.d), which is how they appear on the
// wire. The port is big endian on the wire, unlike every other integer in the
// protocol.
type NetAddress struct {
	// Services is the set of service bits advertised by the peer.
	Services uint64

	// IP is the 16-byte IP address.
	IP [16]byte

	// Port is the TCP port.
	Port uint16
}

// v4MappedPrefix is the 12-byte prefix of an IPv4-mapped IPv6 address.
var v4MappedPrefix = [12]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff}

// NewNetAddress parses host, which must be an IPv4 dotted quad or an IPv6
// literal as defined by RFC 4291 (optionally enclosed in brackets), and
// returns an address with the given port. IPv6 zone identifiers are not
// supported, and DNS names are not resolved: this package performs no name
// service work and does not import the net package.
func NewNetAddress(host string, port uint16) (*NetAddress, error) {
	ip, err := parseIP(host)
	if err != nil {
		return nil, err
	}
	return &NetAddress{IP: ip, Port: port}, nil
}

// IsZero reports whether the address holds no endpoint, that is, whether both
// its IP and its port are zero.
func (na *NetAddress) IsZero() bool {
	if na.Port != 0 {
		return false
	}
	return na.IP == [16]byte{}
}

// IPString returns the textual form of the address's IP: dotted quad for
// IPv4 and IPv4-mapped addresses, and RFC 5952 IPv6 notation otherwise.
func (na *NetAddress) IPString() string {
	return formatIP(na.IP)
}

// String returns the address as host:port, with the host bracketed when it is
// an IPv6 address.
func (na *NetAddress) String() string {
	host := formatIP(na.IP)
	port := strconv.FormatUint(uint64(na.Port), 10)
	if strings.IndexByte(host, ':') >= 0 {
		return "[" + host + "]:" + port
	}
	return host + ":" + port
}

// encode appends the wire encoding of the address: services, IP, and big
// endian port.
func (na *NetAddress) encode(e *Encoder) {
	e.Uint64(na.Services)
	e.Write(na.IP[:])
	e.Uint16BE(na.Port)
}

// decode replaces the address with the one decoded from d.
func (na *NetAddress) decode(d *Decoder) error {
	services, err := d.Uint64()
	if err != nil {
		return fmt.Errorf("btcwire: decoding address services: %w", err)
	}
	ip, err := d.Bytes(len(na.IP))
	if err != nil {
		return fmt.Errorf("btcwire: decoding address IP: %w", err)
	}
	port, err := d.Uint16BE()
	if err != nil {
		return fmt.Errorf("btcwire: decoding address port: %w", err)
	}
	na.Services = services
	copy(na.IP[:], ip)
	na.Port = port
	return nil
}

// parseIP parses host, which may be an IPv4 or IPv6 literal and may be
// bracketed, into its 16-byte wire representation.
func parseIP(host string) ([16]byte, error) {
	var ip [16]byte
	if len(host) >= 2 && host[0] == '[' && host[len(host)-1] == ']' {
		host = host[1 : len(host)-1]
	}
	if host == "" {
		return ip, fmt.Errorf("%w: empty host", ErrInvalidAddress)
	}
	if strings.IndexByte(host, ':') >= 0 {
		parsed, ok := parseIPv6(host)
		if !ok {
			return ip, fmt.Errorf("%w: %q is not an IPv6 address", ErrInvalidAddress, host)
		}
		return parsed, nil
	}
	v4, ok := parseIPv4(host)
	if !ok {
		return ip, fmt.Errorf("%w: %q is not an IPv4 address", ErrInvalidAddress, host)
	}
	copy(ip[:len(v4MappedPrefix)], v4MappedPrefix[:])
	copy(ip[12:16], v4[:])
	return ip, nil
}

// parseIPv4 parses a dotted quad. Leading zeros are rejected because they make
// the interpretation ambiguous.
func parseIPv4(s string) ([4]byte, bool) {
	var out [4]byte
	parts := strings.Split(s, ".")
	if len(parts) != 4 {
		return out, false
	}
	for i, part := range parts {
		if len(part) == 0 || len(part) > 3 {
			return out, false
		}
		if len(part) > 1 && part[0] == '0' {
			return out, false
		}
		n := 0
		for j := 0; j < len(part); j++ {
			c := part[j]
			if c < '0' || c > '9' {
				return out, false
			}
			n = n*10 + int(c-'0')
		}
		if n > 255 {
			return out, false
		}
		out[i] = byte(n)
	}
	return out, true
}

// parseIPv6 parses an IPv6 literal, including forms with an embedded IPv4
// dotted quad and at most one "::" compression.
func parseIPv6(s string) ([16]byte, bool) {
	var ip [16]byte
	if idx := strings.Index(s, "::"); idx >= 0 {
		rest := s[idx+2:]
		if strings.Contains(rest, "::") {
			return ip, false
		}
		head, ok := parseIPv6Groups(s[:idx], false)
		if !ok {
			return ip, false
		}
		tail, ok := parseIPv6Groups(rest, true)
		if !ok {
			return ip, false
		}
		if len(head)+len(tail) > 7 {
			return ip, false
		}
		for i, g := range head {
			binary.BigEndian.PutUint16(ip[2*i:], g)
		}
		base := 8 - len(tail)
		for i, g := range tail {
			binary.BigEndian.PutUint16(ip[2*(base+i):], g)
		}
		return ip, true
	}
	groups, ok := parseIPv6Groups(s, true)
	if !ok || len(groups) != 8 {
		return ip, false
	}
	for i, g := range groups {
		binary.BigEndian.PutUint16(ip[2*i:], g)
	}
	return ip, true
}

// parseIPv6Groups parses the colon-separated 16-bit groups of s. When
// allowIPv4 is true the final group may be a dotted quad, which contributes
// two 16-bit groups.
func parseIPv6Groups(s string, allowIPv4 bool) ([]uint16, bool) {
	if s == "" {
		return nil, true
	}
	parts := strings.Split(s, ":")
	groups := make([]uint16, 0, len(parts))
	for i, part := range parts {
		if part == "" {
			return nil, false
		}
		if allowIPv4 && i == len(parts)-1 && strings.IndexByte(part, '.') >= 0 {
			v4, ok := parseIPv4(part)
			if !ok {
				return nil, false
			}
			groups = append(groups,
				binary.BigEndian.Uint16(v4[0:2]),
				binary.BigEndian.Uint16(v4[2:4]),
			)
			continue
		}
		if len(part) > 4 {
			return nil, false
		}
		var group uint16
		for j := 0; j < len(part); j++ {
			c := part[j]
			var digit uint16
			switch {
			case c >= '0' && c <= '9':
				digit = uint16(c - '0')
			case c >= 'a' && c <= 'f':
				digit = uint16(c-'a') + 10
			case c >= 'A' && c <= 'F':
				digit = uint16(c-'A') + 10
			default:
				return nil, false
			}
			group = group<<4 | digit
		}
		groups = append(groups, group)
	}
	return groups, true
}

// formatIP renders a 16-byte IP in its shortest standard textual form.
func formatIP(ip [16]byte) string {
	if v4, ok := ipv4FromMapped(ip); ok {
		return strconv.Itoa(int(v4[0])) + "." + strconv.Itoa(int(v4[1])) + "." +
			strconv.Itoa(int(v4[2])) + "." + strconv.Itoa(int(v4[3]))
	}
	var groups [8]uint16
	for i := range groups {
		groups[i] = binary.BigEndian.Uint16(ip[2*i:])
	}
	// Find the longest run of zero groups; the leftmost run wins ties. Only
	// runs of two or more groups are compressed.
	bestStart, bestLen := -1, 0
	for i := 0; i < len(groups); {
		if groups[i] != 0 {
			i++
			continue
		}
		j := i
		for j < len(groups) && groups[j] == 0 {
			j++
		}
		if j-i > bestLen {
			bestStart, bestLen = i, j-i
		}
		i = j
	}
	if bestLen < 2 {
		bestStart, bestLen = -1, 0
	}
	var b strings.Builder
	wrote := false
	justCompressed := false
	for i := 0; i < len(groups); i++ {
		if i == bestStart {
			b.WriteString("::")
			wrote, justCompressed = true, true
			i += bestLen - 1
			continue
		}
		if wrote && !justCompressed {
			b.WriteByte(':')
		}
		b.WriteString(strconv.FormatUint(uint64(groups[i]), 16))
		wrote, justCompressed = true, false
	}
	return b.String()
}

// ipv4FromMapped returns the dotted-quad bytes of an IPv4-mapped IPv6
// address.
func ipv4FromMapped(ip [16]byte) ([4]byte, bool) {
	var v4 [4]byte
	for i := 0; i < len(v4MappedPrefix); i++ {
		if ip[i] != v4MappedPrefix[i] {
			return v4, false
		}
	}
	copy(v4[:], ip[12:16])
	return v4, true
}