// Speaker (DaemonSet, hostNetwork): anuncia por ARP las IPs de los Services
// LoadBalancer cuyo "líder" es este nodo.
package main

import (
	"context"
	"crypto/sha256"
	"flag"
	"log"
	"net"
	"net/netip"
	"os"
	"os/signal"
	"sort"
	"sync"
	"syscall"
	"time"

	"github.com/mdlayher/arp"
	"github.com/mdlayher/ethernet"
	corev1 "k8s.io/api/core/v1"
	discoveryv1 "k8s.io/api/discovery/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/tools/clientcmd"
)

const LoadBalancerClass = "lbk8s.io/l2"

type announcer struct {
	mu  sync.RWMutex
	ips map[netip.Addr]bool
	cl  *arp.Client
	hw  net.HardwareAddr
}

func main() {
	kubeconfig := flag.String("kubeconfig", "", "kubeconfig (vacío = in-cluster)")
	ifaceName := flag.String("interface", os.Getenv("IFACE"), "interfaz (vacío = autodetectar por IP del nodo)")
	flag.Parse()
	node := os.Getenv("NODE_NAME")
	if node == "" {
		log.Fatal("falta NODE_NAME")
	}

	cs := mustClient(*kubeconfig)
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	iface, err := pickInterface(ctx, cs, node, *ifaceName)
	if err != nil {
		log.Fatal(err)
	}
	cl, err := arp.Dial(iface)
	if err != nil {
		log.Fatalf("no puedo abrir socket ARP en %s (¿CAP_NET_RAW/hostNetwork?): %v", iface.Name, err)
	}
	defer cl.Close()
	an := &announcer{ips: map[netip.Addr]bool{}, cl: cl, hw: iface.HardwareAddr}
	log.Printf("speaker en nodo %s, interfaz %s (%s)", node, iface.Name, iface.HardwareAddr)
	go an.serve(ctx)

	trigger := make(chan struct{}, 1)
	kick := func() {
		select {
		case trigger <- struct{}{}:
		default:
		}
	}
	h := cache.ResourceEventHandlerFuncs{
		AddFunc:    func(any) { kick() },
		UpdateFunc: func(_, _ any) { kick() },
		DeleteFunc: func(any) { kick() },
	}
	f := informers.NewSharedInformerFactory(cs, 0)
	svcs := f.Core().V1().Services()
	slices := f.Discovery().V1().EndpointSlices()
	nodes := f.Core().V1().Nodes()
	for _, i := range []cache.SharedIndexInformer{svcs.Informer(), slices.Informer(), nodes.Informer()} {
		i.AddEventHandler(h)
	}
	f.Start(ctx.Done())
	f.WaitForCacheSync(ctx.Done())

	tick := time.NewTicker(20 * time.Second) // re-anuncio periódico (GARP)
	defer tick.Stop()
	kick()
	for {
		select {
		case <-ctx.Done():
			return
		case <-trigger:
		case <-tick.C:
		}
		sl, _ := svcs.Lister().List(labels.Everything())
		ns, _ := nodes.Lister().List(labels.Everything())
		es, _ := slices.Lister().List(labels.Everything())
		an.set(desired(node, sl, ns, es))
	}
}

// desired calcula las IPs que debe anunciar este nodo.
func desired(me string, svcs []*corev1.Service, nodes []*corev1.Node, slices []*discoveryv1.EndpointSlice) map[netip.Addr]bool {
	ready := map[string]bool{}
	for _, n := range nodes {
		for _, c := range n.Status.Conditions {
			if c.Type == corev1.NodeReady && c.Status == corev1.ConditionTrue {
				ready[n.Name] = true
			}
		}
	}
	// nodos con endpoints listos por servicio
	local := map[string]map[string]bool{} // ns/svc -> nodos
	for _, s := range slices {
		svc := s.Labels[discoveryv1.LabelServiceName]
		if svc == "" {
			continue
		}
		k := s.Namespace + "/" + svc
		for _, e := range s.Endpoints {
			if e.NodeName == nil || (e.Conditions.Ready != nil && !*e.Conditions.Ready) {
				continue
			}
			if local[k] == nil {
				local[k] = map[string]bool{}
			}
			local[k][*e.NodeName] = true
		}
	}

	out := map[netip.Addr]bool{}
	for _, s := range svcs {
		if s.Spec.Type != corev1.ServiceTypeLoadBalancer {
			continue
		}
		if c := s.Spec.LoadBalancerClass; c != nil && *c != LoadBalancerClass {
			continue
		}
		var ip netip.Addr
		found := false
		for _, ing := range s.Status.LoadBalancer.Ingress {
			if a, err := netip.ParseAddr(ing.IP); err == nil && a.Is4() {
				ip, found = a, true
				break
			}
		}
		if !found {
			continue
		}
		var cands []string
		for n := range ready {
			if s.Spec.ExternalTrafficPolicy == corev1.ServiceExternalTrafficPolicyLocal &&
				!local[s.Namespace+"/"+s.Name][n] {
				continue
			}
			cands = append(cands, n)
		}
		if leader(ip, cands) == me {
			out[ip] = true
		}
	}
	return out
}

// leader elige determinísticamente un nodo por IP (hash estable).
func leader(ip netip.Addr, cands []string) string {
	if len(cands) == 0 {
		return ""
	}
	sort.Slice(cands, func(i, j int) bool {
		return score(ip, cands[i]) < score(ip, cands[j])
	})
	return cands[0]
}

func score(ip netip.Addr, node string) string {
	h := sha256.Sum256([]byte(ip.String() + "/" + node))
	return string(h[:])
}

func (a *announcer) set(want map[netip.Addr]bool) {
	a.mu.Lock()
	prev := a.ips
	a.ips = want
	a.mu.Unlock()
	for ip := range want {
		if !prev[ip] {
			log.Printf("anunciando %s", ip)
		}
		a.gratuitous(ip)
	}
	for ip := range prev {
		if !want[ip] {
			log.Printf("dejo de anunciar %s", ip)
		}
	}
}

func (a *announcer) gratuitous(ip netip.Addr) {
	p, err := arp.NewPacket(arp.OperationReply, a.hw, ip, ethernet.Broadcast, ip)
	if err != nil {
		return
	}
	if err := a.cl.WriteTo(p, ethernet.Broadcast); err != nil {
		log.Printf("GARP %s: %v", ip, err)
	}
}

// serve responde a los ARP request de las IPs anunciadas.
func (a *announcer) serve(ctx context.Context) {
	for ctx.Err() == nil {
		_ = a.cl.SetReadDeadline(time.Now().Add(time.Second))
		pkt, _, err := a.cl.Read()
		if err != nil {
			continue // timeout
		}
		if pkt.Operation != arp.OperationRequest {
			continue
		}
		a.mu.RLock()
		mine := a.ips[pkt.TargetIP]
		a.mu.RUnlock()
		if !mine {
			continue
		}
		if err := a.cl.Reply(pkt, a.hw, pkt.TargetIP); err != nil {
			log.Printf("reply ARP: %v", err)
		}
	}
}

func pickInterface(ctx context.Context, cs kubernetes.Interface, node, name string) (*net.Interface, error) {
	if name != "" {
		return net.InterfaceByName(name)
	}
	n, err := cs.CoreV1().Nodes().Get(ctx, node, metav1Get)
	if err != nil {
		return nil, err
	}
	var nodeIP net.IP
	for _, ad := range n.Status.Addresses {
		if ad.Type == corev1.NodeInternalIP {
			nodeIP = net.ParseIP(ad.Address)
			break
		}
	}
	ifs, _ := net.Interfaces()
	for i := range ifs {
		addrs, _ := ifs[i].Addrs()
		for _, ad := range addrs {
			if ipn, ok := ad.(*net.IPNet); ok && ipn.IP.Equal(nodeIP) {
				return &ifs[i], nil
			}
		}
	}
	return nil, os.ErrNotExist
}

func mustClient(kubeconfig string) kubernetes.Interface {
	var cfg *rest.Config
	var err error
	if kubeconfig != "" {
		cfg, err = clientcmd.BuildConfigFromFlags("", kubeconfig)
	} else {
		cfg, err = rest.InClusterConfig()
	}
	if err != nil {
		log.Fatal(err)
	}
	return kubernetes.NewForConfigOrDie(cfg)
}