// Package lb gives LoadBalancer Services an address that a node holds, and
// forwards that address's traffic to the Service's endpoints with the
// client's address intact.
//
// A Service opts in with spec.loadBalancerClass = Class and names its
// address in spec.loadBalancerIP. The node that holds the address (has it on
// one of its interfaces, as published in its AnnAddresses annotation) is the
// address's holder: it alone writes the Service's status (ipMode Proxy, so
// kube-proxy installs no rules for the address) and DNATs the address's
// ports to ready endpoints, without SNAT. Every other node marks the
// connections that the holder forwards to its pods, and routes their replies
// back to the holder over weave, so that the holder's conntrack can undo the
// DNAT and the reply leaves from the address the client connected to.
//
// Compute is the whole policy as a pure function of what the API shows; the
// datapath and the controller only carry it out.
package lb

import (
	"fmt"
	"hash/fnv"
	"net"
	"sort"
	"strings"

	coreapi "k8s.io/api/core/v1"
	discovery "k8s.io/api/discovery/v1"
)

const (
	// Class is the spec.loadBalancerClass a Service sets to be served by this agent.
	Class = "reweave.levk.dev/node-address"
	// AnnAddresses lists the IPv4 addresses a node holds on its host interfaces (comma-separated).
	AnnAddresses = "reweave.levk.dev/addresses"
	// AnnWeave is "<ip>,<mac>" of the node's weave bridge: where replies to its forwards go.
	AnnWeave = "reweave.levk.dev/weave"

	serviceNameLabel = "kubernetes.io/service-name"
)

// Node is what one node publishes about itself.
type Node struct {
	Name      string
	Addresses []net.IP
	WeaveIP   net.IP
	WeaveMAC  net.HardwareAddr
}

// Holds reports whether the node has ip on one of its host interfaces.
func (n Node) Holds(ip net.IP) bool {
	for _, a := range n.Addresses {
		if a.Equal(ip) {
			return true
		}
	}
	return false
}

// ParseNode reads a Node's annotations. A node that has published nothing
// (yet) holds nothing and cannot take replies.
func ParseNode(n *coreapi.Node) (Node, error) {
	out := Node{Name: n.Name}
	if s := n.Annotations[AnnAddresses]; s != "" {
		for _, f := range strings.Split(s, ",") {
			ip := net.ParseIP(strings.TrimSpace(f)).To4()
			if ip == nil {
				return out, fmt.Errorf("node %s: %s: %q is not an IPv4 address", n.Name, AnnAddresses, f)
			}
			out.Addresses = append(out.Addresses, ip)
		}
	}
	if s := n.Annotations[AnnWeave]; s != "" {
		f := strings.Split(s, ",")
		if len(f) != 2 {
			return out, fmt.Errorf("node %s: %s: want \"<ip>,<mac>\", got %q", n.Name, AnnWeave, s)
		}
		ip := net.ParseIP(f[0]).To4()
		mac, err := net.ParseMAC(f[1])
		if ip == nil || err != nil {
			return out, fmt.Errorf("node %s: %s: want \"<ip>,<mac>\", got %q", n.Name, AnnWeave, s)
		}
		out.WeaveIP, out.WeaveMAC = ip, mac
	}
	return out, nil
}

// FormatAddresses is the AnnAddresses value for a set of addresses (sorted, so it only changes when the set does).
func FormatAddresses(ips []net.IP) string {
	s := make([]string, 0, len(ips))
	for _, ip := range ips {
		s = append(s, ip.String())
	}
	sort.Strings(s)
	return strings.Join(s, ",")
}

// Endpoint is one ready backend of one Service port.
type Endpoint struct {
	IP   net.IP
	Port int32
}

func (e Endpoint) String() string { return fmt.Sprintf("%s:%d", e.IP, e.Port) }

// Port is one Service port at the held address.
type Port struct {
	Protocol  coreapi.Protocol
	Port      int32
	Endpoints []Endpoint
}

// Forward is a Service whose address this node holds.
type Forward struct {
	Namespace, Name string
	Address         net.IP
	Ports           []Port
}

// Status is a write to a Service's status.loadBalancer: Address set, or cleared when nil.
type Status struct {
	Namespace, Name string
	Address         net.IP
}

// Holder is another node that forwards held addresses' connections to pods; Mark identifies its connections here.
type Holder struct {
	Node Node
	Mark uint32 // 1..MaxMarks, in the connmark's top byte
}

// MaxMarks is how many holders a node can tell apart (the top byte of the mark, 0 meaning none).
const MaxMarks = 250

// Desired is what one node must do.
type Desired struct {
	Forwards []Forward
	Holders  []Holder
	Status   []Status
	Problems []string // what an operator should hear about (logged and evented by the controller)
}

// eligible reports whether the Service asks for this agent.
func eligible(s *coreapi.Service) bool {
	return s.Spec.Type == coreapi.ServiceTypeLoadBalancer && s.Spec.LoadBalancerClass != nil &&
		*s.Spec.LoadBalancerClass == Class
}

// shown is the address status.loadBalancer currently shows (nil if none or not a single IP).
func shown(s *coreapi.Service) net.IP {
	if len(s.Status.LoadBalancer.Ingress) != 1 {
		return nil
	}
	return net.ParseIP(s.Status.LoadBalancer.Ingress[0].IP).To4()
}

func hasStatus(s *coreapi.Service) bool { return len(s.Status.LoadBalancer.Ingress) != 0 }

// Compute decides what node `self` does, given every node, every Service and every EndpointSlice.
func Compute(self string, nodes []Node, services []*coreapi.Service, slices []*discovery.EndpointSlice) Desired {
	var d Desired
	sortedNodes := append([]Node(nil), nodes...)
	sort.Slice(sortedNodes, func(i, j int) bool { return sortedNodes[i].Name < sortedNodes[j].Name })
	holdersOf := func(ip net.IP) []Node {
		var hs []Node
		for _, n := range sortedNodes {
			if ip != nil && n.Holds(ip) {
				hs = append(hs, n)
			}
		}
		return hs
	}
	// who may write a status that shows `ip`: its single holder, or (nobody holds it) the first node by name
	writerFor := func(ip net.IP) string {
		if hs := holdersOf(ip); len(hs) == 1 {
			return hs[0].Name
		}
		if len(sortedNodes) > 0 {
			return sortedNodes[0].Name
		}
		return ""
	}
	holding := map[string]bool{} // nodes that forward some address

	svcs := append([]*coreapi.Service(nil), services...)
	sort.Slice(svcs, func(i, j int) bool { return key(svcs[i]) < key(svcs[j]) })
	for _, s := range svcs {
		if !eligible(s) {
			// loadBalancerClass is immutable and only allowed on type LoadBalancer, so a
			// Service never leaves the class while keeping a status we wrote
			continue
		}
		want := net.ParseIP(s.Spec.LoadBalancerIP).To4()
		var holder *Node
		switch hs := holdersOf(want); {
		case want == nil:
			d.Problems = append(d.Problems, fmt.Sprintf("%s: spec.loadBalancerIP must name an IPv4 address a node holds (got %q)", key(s), s.Spec.LoadBalancerIP))
		case len(hs) == 0:
			d.Problems = append(d.Problems, fmt.Sprintf("%s: no node holds %s", key(s), want))
		case len(hs) > 1:
			names := make([]string, len(hs))
			for i, h := range hs {
				names[i] = h.Name
			}
			d.Problems = append(d.Problems, fmt.Sprintf("%s: %s is held by several nodes (%s); serving it from none", key(s), want, strings.Join(names, ", ")))
		default:
			holder = &hs[0]
		}
		if holder == nil {
			// fail closed: nobody forwards, and the status must not claim an address
			if hasStatus(s) && writerFor(shown(s)) == self {
				d.Status = append(d.Status, Status{Namespace: s.Namespace, Name: s.Name})
			}
			continue
		}
		holding[holder.Name] = true
		if holder.Name != self {
			continue
		}
		if cur := shown(s); cur == nil || !cur.Equal(want) {
			d.Status = append(d.Status, Status{Namespace: s.Namespace, Name: s.Name, Address: want})
		}
		d.Forwards = append(d.Forwards, Forward{Namespace: s.Namespace, Name: s.Name, Address: want, Ports: ports(s, slices)})
	}

	var others []Node
	for _, n := range sortedNodes {
		if holding[n.Name] && n.Name != self {
			if n.WeaveIP == nil {
				d.Problems = append(d.Problems, fmt.Sprintf("node %s holds an address but has not published its weave bridge; replies to its forwards cannot reach it", n.Name))
				continue
			}
			others = append(others, n)
		}
	}
	d.Holders = assignMarks(others)
	return d
}

// assignMarks gives each holder a mark from its name, probing past collisions in name order: stable while the set is.
func assignMarks(nodes []Node) []Holder {
	taken := map[uint32]bool{}
	out := make([]Holder, 0, len(nodes))
	for _, n := range nodes {
		if len(out) == MaxMarks {
			break
		}
		h := fnv.New32a()
		h.Write([]byte(n.Name))
		m := h.Sum32()%MaxMarks + 1
		for taken[m] {
			m = m%MaxMarks + 1
		}
		taken[m] = true
		out = append(out, Holder{Node: n, Mark: m})
	}
	return out
}

func key(s *coreapi.Service) string { return s.Namespace + "/" + s.Name }

// ports pairs each Service port with its ready endpoints.
func ports(s *coreapi.Service, slices []*discovery.EndpointSlice) []Port {
	var out []Port
	for _, sp := range s.Spec.Ports {
		proto := sp.Protocol
		if proto == "" {
			proto = coreapi.ProtocolTCP
		}
		seen := map[string]bool{}
		var eps []Endpoint
		for _, sl := range slices {
			if sl.Namespace != s.Namespace || sl.Labels[serviceNameLabel] != s.Name || sl.AddressType != discovery.AddressTypeIPv4 {
				continue
			}
			for _, p := range sl.Ports {
				pname, pproto := "", coreapi.ProtocolTCP
				if p.Name != nil {
					pname = *p.Name
				}
				if p.Protocol != nil {
					pproto = *p.Protocol
				}
				if pname != sp.Name || pproto != proto || p.Port == nil {
					continue
				}
				for _, ep := range sl.Endpoints {
					if ep.Conditions.Ready != nil && !*ep.Conditions.Ready {
						continue
					}
					for _, a := range ep.Addresses {
						ip := net.ParseIP(a).To4()
						if ip == nil {
							continue
						}
						e := Endpoint{IP: ip, Port: *p.Port}
						if !seen[e.String()] {
							seen[e.String()] = true
							eps = append(eps, e)
						}
					}
				}
			}
		}
		sort.Slice(eps, func(i, j int) bool { return eps[i].String() < eps[j].String() })
		out = append(out, Port{Protocol: proto, Port: sp.Port, Endpoints: eps})
	}
	return out
}