package lb

import (
	"fmt"
	"net"
	"strings"
	"testing"

	coreapi "k8s.io/api/core/v1"
	discovery "k8s.io/api/discovery/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func mac(s string) net.HardwareAddr {
	m, err := net.ParseMAC(s)
	if err != nil {
		panic(err)
	}
	return m
}

var (
	cp = Node{Name: "cp", Addresses: []net.IP{net.ParseIP("86.107.168.24").To4()},
		WeaveIP: net.ParseIP("10.32.0.1").To4(), WeaveMAC: mac("aa:aa:aa:aa:aa:01")}
	w1 = Node{Name: "w1", Addresses: []net.IP{net.ParseIP("86.107.168.15").To4()},
		WeaveIP: net.ParseIP("10.40.0.0").To4(), WeaveMAC: mac("aa:aa:aa:aa:aa:02")}
	w2 = Node{Name: "w2", Addresses: []net.IP{net.ParseIP("86.107.168.88").To4()},
		WeaveIP: net.ParseIP("10.44.0.0").To4(), WeaveMAC: mac("aa:aa:aa:aa:aa:03")}
	nodes = []Node{w2, cp, w1}
)

func svc(name, ip string, mutate ...func(*coreapi.Service)) *coreapi.Service {
	class := Class
	s := &coreapi.Service{
		ObjectMeta: metav1.ObjectMeta{Namespace: "gitea", Name: name},
		Spec: coreapi.ServiceSpec{
			Type:              coreapi.ServiceTypeLoadBalancer,
			LoadBalancerClass: &class,
			LoadBalancerIP:    ip,
			Ports: []coreapi.ServicePort{
				{Name: "ssh", Protocol: coreapi.ProtocolTCP, Port: 22},
				{Name: "https", Port: 443}, // protocol defaults to TCP
			},
		},
	}
	for _, m := range mutate {
		m(s)
	}
	return s
}

func shows(ip string) func(*coreapi.Service) {
	return func(s *coreapi.Service) {
		s.Status.LoadBalancer.Ingress = []coreapi.LoadBalancerIngress{{IP: ip}}
	}
}

func ptr[T any](v T) *T { return &v }

func slice(svcName string, eps map[string]bool, ports map[string]int32) *discovery.EndpointSlice {
	sl := &discovery.EndpointSlice{
		ObjectMeta:  metav1.ObjectMeta{Namespace: "gitea", Name: svcName + "-x", Labels: map[string]string{serviceNameLabel: svcName}},
		AddressType: discovery.AddressTypeIPv4,
	}
	for ip, ready := range eps {
		sl.Endpoints = append(sl.Endpoints, discovery.Endpoint{Addresses: []string{ip}, Conditions: discovery.EndpointConditions{Ready: ptr(ready)}})
	}
	for name, port := range ports {
		sl.Ports = append(sl.Ports, discovery.EndpointPort{Name: ptr(name), Protocol: ptr(coreapi.ProtocolTCP), Port: ptr(port)})
	}
	return sl
}

func TestHolderForwardsToReadyEndpointsAndWritesStatus(t *testing.T) {
	s := svc("gitea", "86.107.168.15")
	sl := slice("gitea", map[string]bool{"10.40.0.7": true, "10.44.0.3": true, "10.44.0.9": false},
		map[string]int32{"ssh": 2222, "https": 3000})
	d := Compute("w1", nodes, []*coreapi.Service{s}, []*discovery.EndpointSlice{sl})
	if len(d.Problems) != 0 {
		t.Fatalf("problems: %v", d.Problems)
	}
	if len(d.Forwards) != 1 || !d.Forwards[0].Address.Equal(net.ParseIP("86.107.168.15")) {
		t.Fatalf("forwards: %+v", d.Forwards)
	}
	f := d.Forwards[0]
	if len(f.Ports) != 2 || f.Ports[0].Port != 22 || f.Ports[1].Port != 443 || f.Ports[1].Protocol != coreapi.ProtocolTCP {
		t.Fatalf("ports: %+v", f.Ports)
	}
	if got := f.Ports[0].Endpoints; len(got) != 2 || got[0].String() != "10.40.0.7:2222" || got[1].String() != "10.44.0.3:2222" {
		t.Fatalf("ssh endpoints (not-ready one must be left out): %v", got)
	}
	if got := f.Ports[1].Endpoints; len(got) != 2 || got[0].Port != 3000 {
		t.Fatalf("https endpoints: %v", got)
	}
	if len(d.Status) != 1 || !d.Status[0].Address.Equal(net.ParseIP("86.107.168.15")) {
		t.Fatalf("status: %+v", d.Status)
	}
	if len(d.Holders) != 0 {
		t.Fatalf("a holder does not steer its own replies: %+v", d.Holders)
	}
	// already shown: nothing to write
	d = Compute("w1", nodes, []*coreapi.Service{svc("gitea", "86.107.168.15", shows("86.107.168.15"))}, nil)
	if len(d.Status) != 0 {
		t.Fatalf("status rewritten: %+v", d.Status)
	}
}

func TestOtherNodesSteerRepliesToTheHolderAndNeverWrite(t *testing.T) {
	s := svc("gitea", "86.107.168.15")
	for _, self := range []string{"cp", "w2"} {
		d := Compute(self, nodes, []*coreapi.Service{s}, nil)
		if len(d.Forwards) != 0 || len(d.Status) != 0 {
			t.Fatalf("%s: only the holder forwards and writes: %+v", self, d)
		}
		if len(d.Holders) != 1 || d.Holders[0].Node.Name != "w1" || d.Holders[0].Mark < 1 || d.Holders[0].Mark > MaxMarks {
			t.Fatalf("%s: holders: %+v", self, d.Holders)
		}
	}
}

func TestUnheldOrContestedAddressFailsClosed(t *testing.T) {
	// nobody holds it: no forward anywhere; the stale status is cleared by exactly one node (first by name)
	s := svc("gitea", "86.107.168.99", shows("86.107.168.99"))
	for _, self := range []string{"cp", "w1", "w2"} {
		d := Compute(self, nodes, []*coreapi.Service{s}, nil)
		if len(d.Forwards) != 0 || len(d.Holders) != 0 || len(d.Problems) != 1 || !strings.Contains(d.Problems[0], "no node holds") {
			t.Fatalf("%s: %+v", self, d)
		}
		clears := len(d.Status) == 1 && d.Status[0].Address == nil
		if clears != (self == "cp") {
			t.Fatalf("%s: clears=%v", self, clears)
		}
	}
	// two nodes claim it: neither serves it, its holder's view clears the status
	both := w2
	both.Addresses = append([]net.IP{net.ParseIP("86.107.168.15").To4()}, w2.Addresses...)
	d := Compute("w1", []Node{cp, w1, both}, []*coreapi.Service{svc("gitea", "86.107.168.15", shows("86.107.168.15"))}, nil)
	if len(d.Forwards) != 0 || len(d.Problems) != 1 || !strings.Contains(d.Problems[0], "several nodes") {
		t.Fatalf("contested: %+v", d)
	}
	// a bad request is reported, not guessed at
	d = Compute("w1", nodes, []*coreapi.Service{svc("gitea", "")}, nil)
	if len(d.Forwards) != 0 || len(d.Problems) != 1 || !strings.Contains(d.Problems[0], "spec.loadBalancerIP") {
		t.Fatalf("no address: %+v", d)
	}
}

func TestOnlyOptedInLoadBalancersAreServed(t *testing.T) {
	other := svc("a", "86.107.168.15", func(s *coreapi.Service) { s.Spec.LoadBalancerClass = ptr("example.com/other") })
	noClass := svc("b", "86.107.168.15", func(s *coreapi.Service) { s.Spec.LoadBalancerClass = nil })
	clusterIP := svc("c", "86.107.168.15", func(s *coreapi.Service) {
		s.Spec.Type = coreapi.ServiceTypeClusterIP
	})
	d := Compute("w1", nodes, []*coreapi.Service{other, noClass, clusterIP}, nil)
	if len(d.Forwards)+len(d.Status)+len(d.Holders)+len(d.Problems) != 0 {
		t.Fatalf("touched a Service that did not opt in: %+v", d)
	}
}

func TestParseNodeAndAnnotationsRoundTrip(t *testing.T) {
	l := Local{Addresses: []net.IP{net.ParseIP("86.107.168.88").To4(), net.ParseIP("86.107.168.15").To4()},
		WeaveIP: net.ParseIP("10.44.0.0").To4(), WeaveMAC: mac("aa:aa:aa:aa:aa:03")}
	a := l.Annotations()
	if a[AnnAddresses] != "86.107.168.15,86.107.168.88" || a[AnnWeave] != "10.44.0.0,aa:aa:aa:aa:aa:03" {
		t.Fatalf("annotations: %v", a)
	}
	n, err := ParseNode(&coreapi.Node{ObjectMeta: metav1.ObjectMeta{Name: "w2", Annotations: a}})
	if err != nil || !n.Holds(net.ParseIP("86.107.168.15")) || !n.WeaveIP.Equal(l.WeaveIP) || n.WeaveMAC.String() != "aa:aa:aa:aa:aa:03" {
		t.Fatalf("parsed %+v, %v", n, err)
	}
	for _, bad := range []map[string]string{{AnnAddresses: "1.2.3.4,nope"}, {AnnWeave: "10.44.0.0"}, {AnnWeave: "x,aa:aa:aa:aa:aa:03"}} {
		if _, err := ParseNode(&coreapi.Node{ObjectMeta: metav1.ObjectMeta{Name: "w2", Annotations: bad}}); err == nil {
			t.Fatalf("accepted %v", bad)
		}
	}
	if n, err := ParseNode(&coreapi.Node{ObjectMeta: metav1.ObjectMeta{Name: "new"}}); err != nil || len(n.Addresses) != 0 || n.WeaveIP != nil {
		t.Fatalf("a node that published nothing holds nothing: %+v %v", n, err)
	}
}

func TestMarksAreDistinctAndStable(t *testing.T) {
	var ns []Node
	for _, name := range []string{"a", "b", "c", "d", "e", "f", "g", "h"} {
		ns = append(ns, Node{Name: name})
	}
	first := assignMarks(ns)
	seen := map[uint32]bool{}
	for _, h := range first {
		if seen[h.Mark] || h.Mark < 1 || h.Mark > MaxMarks {
			t.Fatalf("bad or repeated mark: %+v", first)
		}
		seen[h.Mark] = true
	}
	again := assignMarks(ns)
	for i := range first {
		if first[i].Mark != again[i].Mark {
			t.Fatalf("marks moved between passes")
		}
	}
}

func TestRenderedRules(t *testing.T) {
	_, podNet, _ := net.ParseCIDR("10.32.0.0/12")
	s := svc("gitea", "86.107.168.15")
	sl := slice("gitea", map[string]bool{"10.40.0.7": true, "10.44.0.3": true, "10.44.0.4": true}, map[string]int32{"ssh": 2222, "https": 3000})
	d := Compute("w1", nodes, []*coreapi.Service{s}, []*discovery.EndpointSlice{sl})
	r := Render(d, EgressDesired{}, "w1", podNet, []string{"REWEAVE-LBS-DEADBEEF0000"})
	for _, want := range []string{
		"*nat\n:REWEAVE-LB - [0:0]\n:REWEAVE-LB-POST - [0:0]\n:REWEAVE-LB-CLASS - [0:0]\n",
		// the client's address survives weave's masquerade; a pod's own connection (hairpin) is masqueraded here
		"-A REWEAVE-LB-POST -m conntrack --ctstate DNAT --ctorigdst 86.107.168.15/32 -s 10.32.0.0/12 -j MASQUERADE\n",
		"-A REWEAVE-LB-POST -m conntrack --ctstate DNAT --ctorigdst 86.107.168.15/32 -j ACCEPT\n",
		":REWEAVE-LBS-DEADBEEF0000 - [0:0]\n", // stale chain emptied
		"-A REWEAVE-LB -d 86.107.168.15/32 -p tcp -m tcp --dport 22 -m comment --comment \"gitea/gitea\" -j REWEAVE-LBS-",
		"-p tcp -m statistic --mode random --probability 0.3333333333 -j DNAT --to-destination 10.40.0.7:2222\n",
		"-p tcp -m statistic --mode random --probability 0.5000000000 -j DNAT --to-destination 10.44.0.3:2222\n",
		"-p tcp -j DNAT --to-destination 10.44.0.4:2222\n",
		"COMMIT\n",
	} {
		if !strings.Contains(r.Nat, want) {
			t.Fatalf("nat lacks %q:\n%s", want, r.Nat)
		}
	}
	for _, line := range strings.Split(r.Nat, "\n") {
		// iptables refuses a DNAT to ip:port without the protocol in the same rule
		if strings.Contains(line, "-j DNAT") && !strings.Contains(line, " -p tcp ") {
			t.Fatalf("DNAT without its protocol: %q", line)
		}
	}
	if strings.Count(r.Nat, "MASQUERADE") != 1 || strings.Contains(r.Nat, "SNAT") {
		t.Fatalf("only a pod's hairpin connection may be masqueraded:\n%s", r.Nat)
	}
	if r.Mangle != "*mangle\n:REWEAVE-LB-MARK - [0:0]\n:REWEAVE-LB-MSS - [0:0]\nCOMMIT\n" {
		t.Fatalf("a holder marks nothing:\n%s", r.Mangle)
	}
	d = Compute("w2", nodes, []*coreapi.Service{s}, []*discovery.EndpointSlice{sl})
	r = Render(d, EgressDesired{}, "w2", podNet, nil)
	m := d.Holders[0].Mark << 24
	for _, want := range []string{
		"-A REWEAVE-LB-MARK -m conntrack --ctdir REPLY -m connmark ! --mark 0x0/0xff000000 -j CONNMARK --restore-mark --nfmask 0xff000000 --ctmask 0xff000000\n",
		"-A REWEAVE-LB-MARK -i weave -m mac --mac-source aa:aa:aa:aa:aa:02 ! -s 10.32.0.0/12 -d 10.32.0.0/12 -m conntrack --ctstate NEW -m comment --comment \"holder w1\" -j CONNMARK --set-xmark " + fmt.Sprintf("0x%x", m) + "/0xff000000\n",
	} {
		if !strings.Contains(r.Mangle, want) {
			t.Fatalf("mangle lacks %q:\n%s", want, r.Mangle)
		}
	}
	if r.Nat != "*nat\n:REWEAVE-LB - [0:0]\n:REWEAVE-LB-POST - [0:0]\n:REWEAVE-LB-CLASS - [0:0]\nCOMMIT\n" {
		t.Fatalf("a non-holder forwards nothing:\n%s", r.Nat)
	}
	if r.Filter != "*filter\n:REWEAVE-LB-FWD - [0:0]\n:REWEAVE-LB-IN - [0:0]\nCOMMIT\n" {
		t.Fatalf("no EgressService, no filter rules:\n%s", r.Filter)
	}
}