// Package embed implements ip-algo: deterministic embedding of pod identity
// (namespace, pod name, image digest) into the host portion of an IPv6
// address. The mapping is operator-friendly cosmetics — NOT a security
// boundary. See dfritz-cni.md "IPv6 IID Embedding" for the full spec.
package embed

import (
	"encoding/hex"
	"fmt"
	"hash/fnv"
	"net"
	"strings"
)

// Field is one of the supported identity fields.
type Field string

const (
	FieldNamespace Field = "namespace"
	FieldPod       Field = "pod"
	FieldImage     Field = "image"
)

// Values carries the inputs for one embedding call. Image holds the SHA-256
// manifest digest as 64 hex chars when known; otherwise pass the containerID
// in ImageFallback and we'll FNV-1a-64 it.
type Values struct {
	Namespace       string
	Pod             string
	Image           string // 64-char hex sha256 manifest digest, or empty
	ImageFallback   string // typically containerID, used when Image=="".
}

// MaxFieldNibbles is the largest single-field width supported by this
// implementation. 16 nibbles = 64 bits = the output width of FNV-1a-64.
// Wider fields would require a wider hash; the design doc tolerates this
// because real deployments use /64 nodes (15 field nibbles total).
const MaxFieldNibbles = 16

// Embed returns the IPv6 address inside `network` whose host portion encodes
// `fields` (in the given order) followed by the random nibble nNibble.
//
// `network` must be an IPv6 prefix whose length is a multiple of 4 (so the
// host portion is a whole number of nibbles).
//
// `fields` must be non-empty. For a fully-random IID, the caller should pick
// random bytes directly rather than calling Embed.
//
// nNibble is the random "instance" nibble; only the low 4 bits are used.
// Callers regenerate it on collision (see allocations.json).
func Embed(network *net.IPNet, fields []Field, vals Values, nNibble byte) (net.IP, error) {
	ones, bits := network.Mask.Size()
	if bits != 128 {
		return nil, fmt.Errorf("network is not IPv6: %s", network)
	}
	if ones%4 != 0 {
		return nil, fmt.Errorf("prefix length %d is not a multiple of 4", ones)
	}
	hostNibbles := (128 - ones) / 4
	if hostNibbles < 2 {
		return nil, fmt.Errorf("prefix /%d leaves %d host nibble(s); need at least 2 (one field + N)", ones, hostNibbles)
	}
	if len(fields) == 0 {
		return nil, fmt.Errorf("no fields specified; caller should generate random IID directly")
	}

	fieldNibbles := hostNibbles - 1
	dist, err := distribute(fieldNibbles, len(fields))
	if err != nil {
		return nil, err
	}

	addr := make(net.IP, net.IPv6len)
	copy(addr, network.IP.To16())

	// Stream nibbles left-to-right starting at the first host nibble.
	startNibble := ones / 4
	pos := 0
	for i, f := range fields {
		n := dist[i]
		v, err := fieldValue(f, vals, n*4)
		if err != nil {
			return nil, err
		}
		// Write `n` nibbles, most-significant first.
		for j := n - 1; j >= 0; j-- {
			nb := byte((v >> uint(j*4)) & 0xF)
			writeNibble(addr, startNibble+pos, nb)
			pos++
		}
	}
	writeNibble(addr, startNibble+pos, nNibble&0x0F)

	return addr, nil
}

// distribute splits `total` nibbles across `k` fields as evenly as possible,
// giving any remainder to earlier fields one extra nibble at a time.
func distribute(total, k int) ([]int, error) {
	if k <= 0 {
		return nil, fmt.Errorf("k must be > 0")
	}
	if total < k {
		return nil, fmt.Errorf("not enough host nibbles (%d) for %d fields", total, k)
	}
	out := make([]int, k)
	base := total / k
	rem := total % k
	for i := range out {
		out[i] = base
		if i < rem {
			out[i]++
		}
		if out[i] > MaxFieldNibbles {
			return nil, fmt.Errorf("field %d would need %d nibbles; max supported is %d", i, out[i], MaxFieldNibbles)
		}
	}
	return out, nil
}

// fieldValue returns the top `bits` bits of the hash-or-digest for `f`,
// right-aligned in the returned uint64.
func fieldValue(f Field, v Values, bits int) (uint64, error) {
	if bits <= 0 || bits > 64 {
		return 0, fmt.Errorf("bad field bits %d (1..64)", bits)
	}
	switch f {
	case FieldNamespace:
		return topBitsFNV(v.Namespace, bits), nil
	case FieldPod:
		return topBitsFNV(v.Pod, bits), nil
	case FieldImage:
		if v.Image != "" {
			return topBitsHex(v.Image, bits)
		}
		return topBitsFNV(v.ImageFallback, bits), nil
	default:
		return 0, fmt.Errorf("unknown field %q", f)
	}
}

func topBitsFNV(s string, bits int) uint64 {
	h := fnv.New64a()
	_, _ = h.Write([]byte(s))
	return h.Sum64() >> uint(64-bits)
}

// topBitsHex parses a leading sha256-digest-style hex string and returns
// its top `bits` bits, right-aligned. Accepts an optional "sha256:" prefix.
func topBitsHex(s string, bits int) (uint64, error) {
	s = strings.TrimPrefix(s, "sha256:")
	if len(s) < 16 { // need at least 8 bytes / 64 bits to right-shift
		return 0, fmt.Errorf("image digest too short: %d hex chars", len(s))
	}
	b, err := hex.DecodeString(s[:16])
	if err != nil {
		return 0, fmt.Errorf("image digest not hex: %w", err)
	}
	var v uint64
	for _, x := range b {
		v = (v << 8) | uint64(x)
	}
	return v >> uint(64-bits), nil
}

// writeNibble sets the (nibIdx)-th nibble of addr (0 = highest nibble of byte 0).
func writeNibble(addr net.IP, nibIdx int, nb byte) {
	bytePos := nibIdx / 2
	if nibIdx%2 == 0 {
		addr[bytePos] = (addr[bytePos] & 0x0F) | (nb << 4)
	} else {
		addr[bytePos] = (addr[bytePos] & 0xF0) | (nb & 0x0F)
	}
}