package main

import (
	"bytes"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func policy() activationConfig {
	return activationConfig{
		publicInterface:          "ens3",
		minimumObservations:      20,
		maximumAddresses:         4096,
		maximumIntervals:         1024,
		maximumLifetime:          24 * time.Hour,
		maximumObservationWindow: 10 * time.Minute,
	}
}

func validBundle(now time.Time) bundle {
	return bundle{
		SchemaVersion: 1,
		GeneratedAt:   now,
		ExpiresAt:     now.Add(24 * time.Hour),
		DeniedIPv4: []denyEntry{{
			Start:              "192.0.2.1",
			End:                "192.0.2.1",
			Reason:             "connection-rate",
			ObservationStarted: now.Add(-10 * time.Minute),
			ObservationEnded:   now,
			ObservationCount:   20,
			ExpiresAt:          now.Add(24 * time.Hour),
		}},
		DeniedIntervalCount: 1,
		DeniedAddressCount:  1,
	}
}

func TestWriteAndReadEmptyBundle(t *testing.T) {
	now := time.Now().UTC().Truncate(time.Second)
	path := filepath.Join(t.TempDir(), "empty", "current.json.gz")
	document := bundle{
		SchemaVersion: 1,
		GeneratedAt:   now,
		ExpiresAt:     now.Add(24 * time.Hour),
		DeniedIPv4:    []denyEntry{},
	}
	if err := writeBundle(path, document); err != nil {
		t.Fatal(err)
	}
	validated, err := readBundle(now, policy(), path)
	if err != nil {
		t.Fatal(err)
	}
	if validated.DeniedIntervalCount != 0 || validated.DeniedAddressCount != 0 || validated.DeniedIPv4 == nil {
		t.Fatal("empty bundle did not retain an explicit empty deny list")
	}
}

func TestRemovedIntervalCount(t *testing.T) {
	now := time.Date(2026, 7, 26, 0, 0, 0, 0, time.UTC)
	previous, err := validateBundle(now, policy(), validBundle(now))
	if err != nil {
		t.Fatal(err)
	}
	currentDocument := validBundle(now)
	currentDocument.DeniedIPv4 = []denyEntry{}
	currentDocument.DeniedIntervalCount = 0
	currentDocument.DeniedAddressCount = 0
	current, err := validateBundle(now, policy(), currentDocument)
	if err != nil {
		t.Fatal(err)
	}
	if removedIntervalCount(previous, current) != 1 {
		t.Fatal("omitted range was not recorded as removed")
	}
}

func TestWriteJSONFileCreatesParentDirectory(t *testing.T) {
	path := filepath.Join(t.TempDir(), "audit", "summary.json")
	if err := writeJSONFile(path, map[string]int{"removed_interval_count": 1}); err != nil {
		t.Fatal(err)
	}
	contents, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	if string(contents) != "{\"removed_interval_count\":1}\n" {
		t.Fatalf("unexpected summary %q", contents)
	}
}

func TestWriteActivePolicyStatePreservesRemoteCurrentActivation(t *testing.T) {
	now := time.Date(2026, 7, 29, 0, 0, 0, 0, time.UTC)
	directory := t.TempDir()
	validated, err := validateBundle(now, policy(), validBundle(now))
	if err != nil {
		t.Fatal(err)
	}
	if err := writeActivePolicyState(directory, now, "current", &validated); err != nil {
		t.Fatal(err)
	}
	cacheActivation := now.Add(time.Minute)
	if err := writeActivePolicyState(directory, cacheActivation, "cache", &validated); err != nil {
		t.Fatal(err)
	}
	state, err := readActivePolicyState(filepath.Join(directory, activePolicyStateFile))
	if err != nil {
		t.Fatal(err)
	}
	if state.Source != "cache" || !state.ActivatedAt.Equal(cacheActivation) || !state.LastRemoteCurrentAt.Equal(now) {
		t.Fatalf("unexpected active state %#v", state)
	}
	if !state.BundleGeneratedAt.Equal(validated.GeneratedAt) || !state.BundleExpiresAt.Equal(validated.ExpiresAt) {
		t.Fatalf("unexpected bundle state %#v", state)
	}
}

func TestWriteActivePolicyStateRecordsFailOpen(t *testing.T) {
	now := time.Date(2026, 7, 29, 0, 0, 0, 0, time.UTC)
	directory := t.TempDir()
	if err := writeActivePolicyState(directory, now, "fail-open", nil); err != nil {
		t.Fatal(err)
	}
	state, err := readActivePolicyState(filepath.Join(directory, activePolicyStateFile))
	if err != nil {
		t.Fatal(err)
	}
	if state.Source != "fail-open" || !state.BundleGeneratedAt.IsZero() || !state.BundleExpiresAt.IsZero() {
		t.Fatalf("unexpected fail-open state %#v", state)
	}
}

func TestFailOpenRecordsFinalStateAfterRulesLoad(t *testing.T) {
	now := time.Date(2026, 7, 29, 0, 0, 0, 0, time.UTC)
	directory := t.TempDir()
	observedPath := filepath.Join(directory, "observed-state.json")
	nftPath := filepath.Join(directory, "nft")
	script := "#!/bin/sh\nif [ \"$1\" = \"--file\" ]; then cp \"$STATE_PATH\" \"$OBSERVED_PATH\"; exit \"${NFT_LOAD_EXIT:-0}\"; fi\nexit 0\n"
	if err := os.WriteFile(nftPath, []byte(script), 0700); err != nil {
		t.Fatal(err)
	}
	t.Setenv("STATE_PATH", filepath.Join(directory, activePolicyStateFile))
	t.Setenv("OBSERVED_PATH", observedPath)
	configuration := policy()
	configuration.stateDirectory = directory
	configuration.nft = nftPath

	if err := failOpen(now, configuration); err != nil {
		t.Fatal(err)
	}
	observed, err := readActivePolicyState(observedPath)
	if err != nil {
		t.Fatal(err)
	}
	if observed.Source != "transition" {
		t.Fatalf("rules loaded after source %q was reported", observed.Source)
	}
	final, err := readActivePolicyState(filepath.Join(directory, activePolicyStateFile))
	if err != nil {
		t.Fatal(err)
	}
	if final.Source != "fail-open" {
		t.Fatalf("unexpected final source %q", final.Source)
	}
}

func TestActivationRecordsFinalStateAfterRulesLoad(t *testing.T) {
	now := time.Date(2026, 7, 29, 0, 0, 0, 0, time.UTC)
	directory := t.TempDir()
	bundlePath := filepath.Join(directory, "current.json.gz")
	if err := writeBundle(bundlePath, validBundle(now)); err != nil {
		t.Fatal(err)
	}
	observedPath := filepath.Join(directory, "observed-state.json")
	nftPath := filepath.Join(directory, "nft")
	script := "#!/bin/sh\nif [ \"$1\" = \"--file\" ]; then cp \"$STATE_PATH\" \"$OBSERVED_PATH\"; fi\nexit 0\n"
	if err := os.WriteFile(nftPath, []byte(script), 0700); err != nil {
		t.Fatal(err)
	}
	t.Setenv("STATE_PATH", filepath.Join(directory, activePolicyStateFile))
	t.Setenv("OBSERVED_PATH", observedPath)
	configuration := policy()
	configuration.stateDirectory = directory
	configuration.nft = nftPath

	if err := activate(now, configuration, bundlePath, true, "current"); err != nil {
		t.Fatal(err)
	}
	observed, err := readActivePolicyState(observedPath)
	if err != nil {
		t.Fatal(err)
	}
	if observed.Source != "transition" {
		t.Fatalf("rules loaded after source %q was reported", observed.Source)
	}
	final, err := readActivePolicyState(filepath.Join(directory, activePolicyStateFile))
	if err != nil {
		t.Fatal(err)
	}
	if final.Source != "current" {
		t.Fatalf("unexpected final source %q", final.Source)
	}
}

func TestFailOpenLoadFailureRemainsTransitional(t *testing.T) {
	now := time.Date(2026, 7, 29, 0, 0, 0, 0, time.UTC)
	directory := t.TempDir()
	nftPath := filepath.Join(directory, "nft")
	script := "#!/bin/sh\nif [ \"$1\" = \"--file\" ]; then exit 1; fi\nexit 0\n"
	if err := os.WriteFile(nftPath, []byte(script), 0700); err != nil {
		t.Fatal(err)
	}
	configuration := policy()
	configuration.stateDirectory = directory
	configuration.nft = nftPath

	if err := failOpen(now, configuration); err == nil {
		t.Fatal("expected nftables load failure")
	}
	state, err := readActivePolicyState(filepath.Join(directory, activePolicyStateFile))
	if err != nil {
		t.Fatal(err)
	}
	if state.Source != "transition" {
		t.Fatalf("failed rules load reported source %q", state.Source)
	}
}

func TestValidateAndRender(t *testing.T) {
	now := time.Date(2026, 7, 26, 0, 0, 0, 0, time.UTC)
	validated, err := validateBundle(now, policy(), validBundle(now))
	if err != nil {
		t.Fatal(err)
	}
	var output bytes.Buffer
	if err := render(&output, now, validated, "ens3"); err != nil {
		t.Fatal(err)
	}
	for _, expected := range []string{
		"table inet oci_l4_relay_smtp_deny",
		"flags interval, timeout",
		"192.0.2.1 timeout 86400s",
		`iifname "ens3" tcp dport 25`,
		"counter name smtp_denied drop",
	} {
		if !strings.Contains(output.String(), expected) {
			t.Errorf("rendered ruleset does not contain %q", expected)
		}
	}
	for _, excluded := range []string{"dport 465", "dport 587", "dport 993", "dport { 80, 443 }"} {
		if strings.Contains(output.String(), excluded) {
			t.Errorf("rendered SMTP deny ruleset contains %q", excluded)
		}
	}
}

func TestEmptyBundleIsValidAndFailOpenIsEmpty(t *testing.T) {
	now := time.Date(2026, 7, 26, 0, 0, 0, 0, time.UTC)
	document := validBundle(now)
	document.DeniedIPv4 = nil
	document.DeniedIntervalCount = 0
	document.DeniedAddressCount = 0
	validated, err := validateBundle(now, policy(), document)
	if err != nil {
		t.Fatal(err)
	}
	var rendered bytes.Buffer
	if err := render(&rendered, now, validated, "ens3"); err != nil {
		t.Fatal(err)
	}
	var failOpen bytes.Buffer
	if err := renderFailOpen(&failOpen, "ens3"); err != nil {
		t.Fatal(err)
	}
	for name, value := range map[string]string{"empty bundle": rendered.String(), "fail open": failOpen.String()} {
		if strings.Contains(value, "elements =") {
			t.Errorf("%s unexpectedly contains deny elements", name)
		}
		if !strings.Contains(value, "ip saddr @denied_v4") {
			t.Errorf("%s does not retain the empty deny lookup", name)
		}
	}
}

func TestExpiredEntryDecaysWithoutInvalidatingBundle(t *testing.T) {
	now := time.Date(2026, 7, 26, 0, 0, 0, 0, time.UTC)
	document := validBundle(now)
	document.DeniedIPv4[0].ObservationEnded = now.Add(-2 * time.Minute)
	document.DeniedIPv4[0].ExpiresAt = now.Add(-time.Minute)
	validated, err := validateBundle(now, policy(), document)
	if err != nil {
		t.Fatal(err)
	}
	var output bytes.Buffer
	if err := render(&output, now, validated, "ens3"); err != nil {
		t.Fatal(err)
	}
	if strings.Contains(output.String(), "192.0.2.1") {
		t.Fatal("expired entry remained in the rendered deny set")
	}
}

func TestPolicyBoundaries(t *testing.T) {
	now := time.Date(2026, 7, 26, 0, 0, 0, 0, time.UTC)
	maximum := validBundle(now)
	maximum.DeniedIPv4[0].Start = "192.0.2.0"
	maximum.DeniedIPv4[0].End = "192.0.17.255"
	maximum.DeniedAddressCount = 4096
	if _, err := validateBundle(now, policy(), maximum); err != nil {
		t.Fatal(err)
	}
	tests := map[string]func(*bundle){
		"observation threshold": func(value *bundle) {
			value.DeniedIPv4[0].ObservationCount = 19
		},
		"address cap": func(value *bundle) {
			value.DeniedIPv4[0].Start = "192.0.2.0"
			value.DeniedIPv4[0].End = "192.0.18.0"
			value.DeniedAddressCount = 4097
		},
		"lifetime": func(value *bundle) {
			value.ExpiresAt = now.Add(24*time.Hour + time.Second)
			value.DeniedIPv4[0].ExpiresAt = value.ExpiresAt
		},
		"observation window": func(value *bundle) {
			value.DeniedIPv4[0].ObservationStarted = now.Add(-10*time.Minute - time.Second)
		},
	}
	for name, mutate := range tests {
		t.Run(name, func(t *testing.T) {
			value := validBundle(now)
			mutate(&value)
			if _, err := validateBundle(now, policy(), value); err == nil {
				t.Fatal("expected policy-bound failure")
			}
		})
	}
	intervals := validBundle(now)
	intervals.DeniedIPv4 = make([]denyEntry, 1025)
	for index := range intervals.DeniedIPv4 {
		entry := validBundle(now).DeniedIPv4[0]
		entry.Start = ipv4String(uint32(index * 2))
		entry.End = entry.Start
		intervals.DeniedIPv4[index] = entry
	}
	intervals.DeniedIntervalCount = len(intervals.DeniedIPv4)
	intervals.DeniedAddressCount = uint64(len(intervals.DeniedIPv4))
	if _, err := validateBundle(now, policy(), intervals); err == nil {
		t.Fatal("expected interval-cap failure")
	}
}

func TestValidationRejectsNonTransportOrMalformedEvidence(t *testing.T) {
	now := time.Date(2026, 7, 26, 0, 0, 0, 0, time.UTC)
	tests := map[string]func(*bundle){
		"schema":  func(value *bundle) { value.SchemaVersion = 2 },
		"expired": func(value *bundle) { value.ExpiresAt = now },
		"ipv6":    func(value *bundle) { value.DeniedIPv4[0].Start = "2001:db8::1" },
		"reason":  func(value *bundle) { value.DeniedIPv4[0].Reason = "spam-score" },
		"future window": func(value *bundle) {
			value.DeniedIPv4[0].ObservationEnded = now.Add(6 * time.Minute)
		},
	}
	for name, mutate := range tests {
		t.Run(name, func(t *testing.T) {
			value := validBundle(now)
			mutate(&value)
			if _, err := validateBundle(now, policy(), value); err == nil {
				t.Fatal("expected validation failure")
			}
		})
	}
}

func TestReadBundleRejectsContentEvidence(t *testing.T) {
	now := time.Now().UTC().Truncate(time.Second)
	contents, err := json.Marshal(validBundle(now))
	if err != nil {
		t.Fatal(err)
	}
	contents = bytes.Replace(contents, []byte(`"reason":"connection-rate"`), []byte(`"reason":"connection-rate","spam_score":9`), 1)
	path := filepath.Join(t.TempDir(), "bundle.json")
	if err := os.WriteFile(path, contents, 0600); err != nil {
		t.Fatal(err)
	}
	if _, err := readBundle(now, policy(), path); err == nil {
		t.Fatal("expected content evidence to be rejected")
	}
}

func TestRejectsUnsafeInterface(t *testing.T) {
	if err := renderFailOpen(&bytes.Buffer{}, "ens3\nflush ruleset"); err == nil {
		t.Fatal("expected unsafe interface rejection")
	}
}