package main

import (
	"context"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/google/nftables"
	"github.com/google/nftables/expr"
	discoveryv1 "k8s.io/api/discovery/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes/fake"
)

type roundTripperFunc func(*http.Request) (*http.Response, error)

func (function roundTripperFunc) RoundTrip(request *http.Request) (*http.Response, error) {
	return function(request)
}

func TestRepeatElectionAfterTermFailure(t *testing.T) {
	ctx, stop := context.WithCancel(context.Background())
	var attempts atomic.Int32
	repeatElection(ctx, time.Millisecond, 10*time.Millisecond, func(_ context.Context, cancelElection context.CancelFunc) bool {
		if attempts.Add(1) == 1 {
			cancelElection()
			return false
		}
		stop()
		return false
	})
	if attempts.Load() != 2 {
		t.Fatalf("expected two election attempts, got %d", attempts.Load())
	}
}

func TestElectionRetryDelay(t *testing.T) {
	base := 5 * time.Second
	maximum := 60 * time.Second
	for unhealthyTerms, expected := range map[int]time.Duration{
		0:  base,
		1:  base,
		2:  10 * time.Second,
		3:  20 * time.Second,
		4:  40 * time.Second,
		5:  maximum,
		20: maximum,
	} {
		if actual := electionRetryDelay(base, maximum, unhealthyTerms); actual != expected {
			t.Fatalf("terms=%d: got %s, want %s", unhealthyTerms, actual, expected)
		}
	}
}

func TestParseCIDRPreservesHostAddress(t *testing.T) {
	address, err := parseCIDR("10.153.238.225/30")
	if err != nil {
		t.Fatal(err)
	}
	if got := address.String(); got != "10.153.238.225/30" {
		t.Fatalf("got %q", got)
	}
}

func TestParseMetricsTarget(t *testing.T) {
	tunnel, err := parseCIDR("10.153.238.225/30")
	if err != nil {
		t.Fatal(err)
	}
	target, err := parseMetricsTarget("http://10.40.68.115:19100/metrics", tunnel)
	if err != nil {
		t.Fatal(err)
	}
	if target.String() != "http://10.40.68.115:19100/metrics" {
		t.Fatalf("unexpected target %q", target)
	}
	for _, value := range []string{
		"https://10.40.68.115:19100/metrics",
		"http://10.153.238.225:19100/metrics",
		"http://10.40.68.115:19100/other",
		"http://10.40.68.115:19100/metrics?query=true",
		"http://192.0.2.1:19100/metrics",
	} {
		if _, err := parseMetricsTarget(value, tunnel); err == nil {
			t.Fatalf("expected %q to fail", value)
		}
	}
}

func TestMetricsHandlerRequiresActiveOwner(t *testing.T) {
	target, err := url.Parse("http://10.40.68.115:19100/metrics")
	if err != nil {
		t.Fatal(err)
	}
	client := &http.Client{Transport: roundTripperFunc(func(request *http.Request) (*http.Response, error) {
		if request.URL.String() != target.String() {
			t.Fatalf("unexpected upstream request %q", request.URL)
		}
		return &http.Response{
			StatusCode: http.StatusOK,
			Header:     http.Header{"Content-Type": []string{"text/plain; version=0.0.4"}},
			Body:       io.NopCloser(strings.NewReader("relay_metric 1\n")),
		}, nil
	})}
	instance := &relay{config: config{metricsTarget: target}, metricsClient: client}
	inactive := httptest.NewRecorder()
	instance.metricsHandler().ServeHTTP(inactive, httptest.NewRequest(http.MethodGet, "/metrics", nil))
	if inactive.Code != http.StatusServiceUnavailable {
		t.Fatalf("inactive status = %d", inactive.Code)
	}
	instance.active.Store(true)
	active := httptest.NewRecorder()
	instance.metricsHandler().ServeHTTP(active, httptest.NewRequest(http.MethodGet, "/metrics", nil))
	if active.Code != http.StatusOK || active.Body.String() != "relay_metric 1\n" || active.Header().Get("Content-Type") != "text/plain; version=0.0.4" {
		t.Fatalf("unexpected active response: status=%d content-type=%q body=%q", active.Code, active.Header().Get("Content-Type"), active.Body.String())
	}
}

func TestMetricsEndpointFollowsOwner(t *testing.T) {
	client := fake.NewSimpleClientset(&discoveryv1.EndpointSlice{
		ObjectMeta:  metav1.ObjectMeta{Name: "oci-l4-relay-metrics", Namespace: "kube-system"},
		AddressType: discoveryv1.AddressTypeIPv4,
	})
	owner := &relay{
		client: client,
		config: config{
			namespace:        "kube-system",
			podName:          "relay-owner",
			podUID:           "owner-uid",
			podIP:            "10.167.40.23",
			podNodeName:      "node-a",
			metricsSliceName: "oci-l4-relay-metrics",
		},
	}
	if err := owner.publishMetricsEndpoint(context.Background()); err != nil {
		t.Fatal(err)
	}
	slice, err := client.DiscoveryV1().EndpointSlices("kube-system").Get(context.Background(), "oci-l4-relay-metrics", metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if !owner.metricsEndpointMatches(slice) {
		t.Fatalf("endpoint does not match owner: %#v", slice.Endpoints)
	}
	standby := &relay{client: client, config: config{namespace: "kube-system", podUID: "standby-uid", metricsSliceName: "oci-l4-relay-metrics"}}
	if err := standby.clearMetricsEndpoint(context.Background()); err != nil {
		t.Fatal(err)
	}
	slice, err = client.DiscoveryV1().EndpointSlices("kube-system").Get(context.Background(), "oci-l4-relay-metrics", metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if len(slice.Endpoints) != 1 {
		t.Fatal("standby cleared the active owner's endpoint")
	}
	if err := owner.clearMetricsEndpoint(context.Background()); err != nil {
		t.Fatal(err)
	}
	slice, err = client.DiscoveryV1().EndpointSlices("kube-system").Get(context.Background(), "oci-l4-relay-metrics", metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if len(slice.Endpoints) != 0 {
		t.Fatal("owner endpoint was not cleared")
	}
}

func TestAllowedRuleContainsGatewayAndPort(t *testing.T) {
	rule := allowedRule(net.ParseIP("10.104.40.127"), 443)
	if len(rule) != 12 {
		t.Fatalf("got %d expressions", len(rule))
	}
}

func TestMetricsReturnRuleIsBoundedToEstablishedPeerTraffic(t *testing.T) {
	rule := metricsReturnRule(net.ParseIP("10.153.238.225"), net.ParseIP("10.40.68.115"), 19100)
	if len(rule) != 16 {
		t.Fatalf("got %d expressions", len(rule))
	}
	comparison, ok := rule[14].(*expr.Cmp)
	if !ok || comparison.Op != expr.CmpOpNeq {
		t.Fatal("metrics return rule does not require a matching conntrack state")
	}
}

func TestParseRelayTargets(t *testing.T) {
	targets, err := parseRelayTargets(`[{"address":"10.66.153.202","ports":[993,25,465]},{"address":"10.104.40.127","ports":[443,80]}]`)
	if err != nil {
		t.Fatal(err)
	}
	if len(targets) != 2 ||
		targets[0].address.String() != "10.104.40.127" ||
		targets[1].address.String() != "10.66.153.202" ||
		targets[1].ports[0] != 25 ||
		targets[1].ports[2] != 993 {
		t.Fatalf("unexpected normalized targets: %#v", targets)
	}
}

func TestParseRelayTargetsRejectsAmbiguousMappings(t *testing.T) {
	for name, value := range map[string]string{
		"unknown field":  `[{"address":"10.104.40.127","ports":[80],"name":"web"}]`,
		"duplicate port": `[{"address":"10.104.40.127","ports":[80]},{"address":"10.66.153.202","ports":[80]}]`,
		"unknown port":   `[{"address":"10.104.40.127","ports":[22]}]`,
		"submission":     `[{"address":"10.66.153.202","ports":[587]}]`,
		"duplicate IP":   `[{"address":"10.104.40.127","ports":[80]},{"address":"10.104.40.127","ports":[443]}]`,
		"IPv6":           `[{"address":"2001:db8::1","ports":[443]}]`,
		"trailing":       `[{"address":"10.104.40.127","ports":[80]}] true`,
	} {
		t.Run(name, func(t *testing.T) {
			if _, err := parseRelayTargets(value); err == nil {
				t.Fatal("expected target validation failure")
			}
		})
	}
}

func TestParseNetworks(t *testing.T) {
	networks, err := parseNetworks("10.63.81.107/24, 10.162.8.60/24")
	if err != nil {
		t.Fatal(err)
	}
	if len(networks) != 2 || networks[0].String() != "10.63.81.107/24" || networks[1].String() != "10.162.8.60/24" {
		t.Fatalf("unexpected networks: %v", networks)
	}
}

func TestFirewallCountersFromRules(t *testing.T) {
	rule := func(packets uint64) *nftables.Rule {
		return &nftables.Rule{Exprs: []expr.Any{&expr.Counter{Packets: packets}}}
	}
	targets := []relayTarget{{address: net.ParseIP("10.104.40.127").To4(), ports: []uint16{80, 443}}}
	counters, err := firewallCountersFromRules(map[string][]*nftables.Rule{
		"prerouting": {{}, rule(1), rule(2), rule(3)},
		"output":     {rule(4), rule(5)},
	}, targets)
	if err != nil {
		t.Fatal(err)
	}
	if counters.tunnelAllowed["10.104.40.127:80"] != 1 || counters.tunnelAllowed["10.104.40.127:443"] != 2 || counters.tunnelDenied != 3 || counters.outerAllowed != 4 || counters.outerDenied != 5 {
		t.Fatalf("unexpected counters: %+v", counters)
	}
}

func TestFirewallCountersRequireEveryRuleCounter(t *testing.T) {
	withCounter := &nftables.Rule{Exprs: []expr.Any{&expr.Counter{}}}
	withoutCounter := &nftables.Rule{}
	_, err := firewallCountersFromRules(map[string][]*nftables.Rule{
		"prerouting": {{}, withCounter, withoutCounter, withCounter},
		"output":     {withCounter, withCounter},
	}, []relayTarget{{address: net.ParseIP("10.104.40.127").To4(), ports: []uint16{80, 443}}})
	if err == nil {
		t.Fatal("expected missing counter error")
	}
}

func TestCheckHealth(t *testing.T) {
	path := filepath.Join(t.TempDir(), "health")
	if err := writeHealth(path, true); err != nil {
		t.Fatal(err)
	}
	if err := checkHealth(path, time.Minute); err != nil {
		t.Fatal(err)
	}
	if err := writeHealth(path, false); err != nil {
		t.Fatal(err)
	}
	if err := checkHealth(path, time.Minute); err == nil || err.Error() != "relay reports unhealthy" {
		t.Fatalf("expected unhealthy state failure, got %v", err)
	}
	if err := writeHealth(path, true); err != nil {
		t.Fatal(err)
	}
	old := time.Now().Add(-time.Hour)
	if err := os.Chtimes(path, old, old); err != nil {
		t.Fatal(err)
	}
	if err := checkHealth(path, time.Minute); err == nil {
		t.Fatal("expected stale health failure")
	}
}

func TestVerifyReversePathFilterAcceptsEffectiveLooseMode(t *testing.T) {
	for name, values := range map[string]map[string]string{
		"interface": {"all": "0\n", interfaceName: "2\n"},
		"all":       {"all": "2\n", interfaceName: "1\n"},
	} {
		t.Run(name, func(t *testing.T) {
			root := t.TempDir()
			for directoryName, value := range values {
				directory := filepath.Join(root, directoryName)
				if err := os.MkdirAll(directory, 0755); err != nil {
					t.Fatal(err)
				}
				if err := os.WriteFile(filepath.Join(directory, "rp_filter"), []byte(value), 0644); err != nil {
					t.Fatal(err)
				}
			}
			if err := verifyReversePathFilter(root, interfaceName); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestVerifyReversePathFilterRejectsStrictMode(t *testing.T) {
	root := t.TempDir()
	for _, name := range []string{"all", interfaceName} {
		directory := filepath.Join(root, name)
		if err := os.MkdirAll(directory, 0755); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(filepath.Join(directory, "rp_filter"), []byte("1\n"), 0644); err != nil {
			t.Fatal(err)
		}
	}
	if err := verifyReversePathFilter(root, interfaceName); err == nil {
		t.Fatal("expected strict reverse-path filtering to fail")
	}
}