package cli

import (
	"bytes"
	"context"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

type mockNftTransport struct {
	execOut map[string][]byte
	execErr map[string]error
	execs   []string
}

func (m *mockNftTransport) Exec(ctx context.Context, peer string, cmd string) ([]byte, error) {
	m.execs = append(m.execs, cmd)
	if m.execErr != nil {
		if err, ok := m.execErr[cmd]; ok {
			return nil, err
		}
	}
	if m.execOut != nil {
		if out, ok := m.execOut[cmd]; ok {
			return out, nil
		}
	}
	return nil, nil
}

func (m *mockNftTransport) ReadFile(ctx context.Context, peer string, path string) ([]byte, error) {
	return nil, errors.New("not implemented")
}

func TestDoctorNft_AllPass(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	origTransport := nftTransportOverride
	origPeer := nftLeadPeerOverride
	defer func() {
		nftTransportOverride = origTransport
		nftLeadPeerOverride = origPeer
	}()
	nftLeadPeerOverride = "lead"
	nftTransportOverride = &mockNftTransport{
		execOut: map[string][]byte{
			"nft list table inet orca-ingress": []byte(`table inet orca-ingress {
	set orca_trusted_probes { type ipv4_addr; }
	chain prerouting {
		tcp dport 443 dnat to 127.0.0.1:8443
		tcp dport 80 dnat to 127.0.0.1:8080
	}
	chain forward {
		tcp dport 443 ct state new meter { ora_rl { rate 100/second } } accept
	}
}`),
		},
	}
	resetRootFlags(t)
	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"doctor", "nft"})
	if err := rootCmd.Execute(); err != nil {
		t.Fatalf("doctor nft: %v", err)
	}
	out := buf.String()
	for _, want := range []string{"nft:table", "nft:dnat-443", "nft:dnat-80", "nft:rate-limit", "nft:parse", "PASS"} {
		if !strings.Contains(out, want) {
			t.Errorf("output missing %q:\n%s", want, out)
		}
	}
}

func TestDoctorNft_TableMissing(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	origTransport := nftTransportOverride
	origPeer := nftLeadPeerOverride
	defer func() {
		nftTransportOverride = origTransport
		nftLeadPeerOverride = origPeer
	}()
	nftLeadPeerOverride = "lead"
	nftTransportOverride = &mockNftTransport{
		execErr: map[string]error{
			"nft list table inet orca-ingress": errors.New("table not found"),
		},
	}
	resetRootFlags(t)
	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"doctor", "nft"})
	if err := rootCmd.Execute(); err != nil {
		t.Fatalf("doctor nft: %v", err)
	}
	out := buf.String()
	if !strings.Contains(out, "nft:table") || !strings.Contains(out, "FAIL") {
		t.Errorf("expected table FAIL:\n%s", out)
	}
}

func TestDoctorNft_HashDriftDetected(t *testing.T) {
	dir, cleanup := initTestEnv(t)
	defer cleanup()
	origTransport := nftTransportOverride
	origPeer := nftLeadPeerOverride
	defer func() {
		nftTransportOverride = origTransport
		nftLeadPeerOverride = origPeer
	}()
	nftLeadPeerOverride = "lead"
	rendered := `table inet orca-ingress { tcp dport 443 dnat to 127.0.0.1:8443; tcp dport 80 dnat to 127.0.0.1:8080; ora_rl; }`
	nftTransportOverride = &mockNftTransport{
		execOut: map[string][]byte{
			"nft list table inet orca-ingress":                       []byte(rendered),
			"nft -c -f /etc/nftables.d/orca.nft":                     []byte(""),
			"sha256sum /etc/nftables.d/orca.nft 2>/dev/null":          []byte("deadbeef  /etc/nftables.d/orca.nft\n"),
		},
	}
	// Record a DIFFERENT hash so drift is reported.
	if err := os.MkdirAll(filepath.Dir(filepath.Join(dir, "cluster", "nft.applied.sha256")), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "cluster", "nft.applied.sha256"), []byte("cafef00d\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	resetRootFlags(t)
	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"doctor", "nft"})
	if err := rootCmd.Execute(); err != nil {
		t.Fatalf("doctor nft: %v", err)
	}
	out := buf.String()
	if !strings.Contains(out, "nft:hash-drift") || !strings.Contains(out, "DRIFT") || !strings.Contains(out, "FAIL") {
		t.Errorf("expected DRIFT FAIL:\n%s", out)
	}
}
