package nftables

import (
	"testing"

	"github.com/google/nftables/expr"
)

func TestRestoreChangeSet(t *testing.T) {
	accept := []expr.Any{&expr.Verdict{Kind: expr.VerdictAccept}}
	drop := []expr.Any{&expr.Verdict{Kind: expr.VerdictDrop}}

	snap := &FirewallState{Rules: map[string][]ManagedRule{
		"input": {
			{Chain: "input", Tag: "ssh", Exprs: accept, Handle: 4},
			{Chain: "input", Tag: "web", Exprs: accept, Handle: 5},
			{Chain: "input", Tag: "", Exprs: drop, Handle: 6},
		},
		"forward": {{Chain: "forward", Tag: "fwd", Exprs: accept, Handle: 7}},
	}}
	current := &FirewallState{Rules: map[string][]ManagedRule{
		"input": {
			{Chain: "input", Tag: "web", Exprs: accept, Handle: 10},
			{Chain: "input", Tag: "ssh", Exprs: drop, Handle: 11},
			{Chain: "input", Tag: "", Exprs: drop, Handle: 12},
		},
	}}

	cs := restoreChangeSet(current, snap)

	var removed []uint64
	for _, r := range cs.Remove {
		removed = append(removed, r.Handle)
	}
	if len(removed) != 2 || removed[0] != 10 || removed[1] != 11 {
		t.Errorf("expected managed handles [10 11] removed, untagged kept; got %v", removed)
	}

	var added []string
	for _, r := range cs.Add {
		added = append(added, r.Tag)
	}
	want := []string{"fwd", "ssh", "web"}
	if len(added) != len(want) {
		t.Fatalf("added %v, want %v", added, want)
	}
	for i := range want {
		if added[i] != want[i] {
			t.Fatalf("added %v, want %v (snapshot order per chain)", added, want)
		}
	}
	if !exprsEqual(cs.Add[1].Exprs, accept) {
		t.Error("ssh not restored to its snapshot exprs")
	}
}

func TestRestoreChangeSetEmptySnapshotRemovesAll(t *testing.T) {
	current := &FirewallState{Rules: map[string][]ManagedRule{
		"input": {{Chain: "input", Tag: "x", Handle: 1}},
	}}
	cs := restoreChangeSet(current, &FirewallState{Rules: map[string][]ManagedRule{}})
	if len(cs.Remove) != 1 || len(cs.Add) != 0 {
		t.Errorf("expected 1 remove 0 add, got %d/%d", len(cs.Remove), len(cs.Add))
	}
}