package cli

import (
	"context"
	"net/netip"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/firewall"
)

type scriptLog struct {
	scripts []string
	reply   func(script string) string
}

func (s *scriptLog) Run(_ context.Context, script string, _ []byte) ([]byte, error) {
	s.scripts = append(s.scripts, script)
	if s.reply != nil {
		return []byte(s.reply(script)), nil
	}
	return nil, nil
}

func TestAControlPlaneFollowsTheFirewallSetting(t *testing.T) {
	off := &scriptLog{}
	if _, err := syncFirewall(context.Background(), off, firewallState{Enabled: false}); err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(strings.Join(off.scripts, "\n"), "nft delete table inet llmgw") {
		t.Errorf("off did not take the table out: %q", off.scripts)
	}

	none := &scriptLog{reply: func(s string) string {
		if strings.Contains(s, "echo loaded") {
			return "none\n"
		}
		return ""
	}}
	msg, err := syncFirewall(context.Background(), none, firewallState{Enabled: true})
	if err != nil || !strings.Contains(msg, "server init") {
		t.Errorf("on, with no rules here: %q %v", msg, err)
	}

	on := &scriptLog{reply: func(s string) string {
		if strings.Contains(s, "echo loaded") {
			return "loaded\n"
		}
		return ""
	}}
	if msg, err := syncFirewall(context.Background(), on, firewallState{Enabled: true}); err != nil || msg != "llmgw's firewall is on here" {
		t.Errorf("on: %q %v", msg, err)
	}
}

func TestAMachineLetsInTheClusterMachinesItIsTold(t *testing.T) {
	// The rules hold 10.0.0.0/24; a node on 192.168.7.0/24 is let in, as
	// the route to it says, and one already in is not asked about.
	r := &scriptLog{reply: func(s string) string {
		switch {
		case strings.Contains(s, "echo loaded"):
			return "loaded\n"
		case strings.Contains(s, "cat /etc/llmgw/firewall.nft"):
			return firewallRulesFor("10.0.0.0/24")
		case strings.Contains(s, "ip -o route get 192.168.7.20"):
			return "192.168.7.20 dev eth1 src 192.168.7.1 uid 0\n2: eth1    inet 192.168.7.1/24 brd 192.168.7.255 scope global eth1\n"
		case strings.Contains(s, "nft -f"):
			return "changed\n"
		}
		return ""
	}}
	msg, err := syncFirewall(context.Background(), r, firewallState{Enabled: true, Toward: []string{"10.0.0.147", "192.168.7.20"}})
	if err != nil || !strings.Contains(msg, "192.168.7.0/24") {
		t.Fatalf("%q %v", msg, err)
	}
	if strings.Contains(strings.Join(r.scripts, "\n"), "route get 10.0.0.147") {
		t.Error("asked the route to a machine already let in")
	}
}

func firewallRulesFor(nets string) string {
	return firewall.Script(firewall.Spec{Ports: []firewall.Port{firewall.MirrorFiles}, Allow: []netip.Prefix{netip.MustParsePrefix(nets)}})
}