//go:build integration

package wginterface_test

import (
	"fmt"
	"log"
	"net"
	"net/netip"
	"os"
	"strings"
	"testing"
	"time"

	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"vppn/peer/wginterface"
)

const (
	testBasePort = 59100
)

func TestMain(m *testing.M) {
	if os.Getuid() != 0 {
		fmt.Fprintln(os.Stderr, "wginterface integration tests require root; skipping")
		os.Exit(0)
	}
	os.Exit(m.Run())
}

type testPeer struct {
	Name    string
	VpnIP   netip.Addr
	Port    int
	PrivKey wgtypes.Key
	PubKey  wgtypes.Key
	Dev     *wginterface.Device
}

func (p *testPeer) Endpoint() netip.AddrPort {
	return netip.AddrPortFrom(netip.AddrFrom4([4]byte{127, 0, 0, 1}), uint16(p.Port))
}

func newTestPeer(t *testing.T, name string, vpnIP netip.Addr, port int) *testPeer {
	t.Helper()

	privKey, err := wgtypes.GenerateKey()
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}

	a4 := vpnIP.As4()
	if err := wginterface.Create(name, net.IP(a4[:]), 24); err != nil {
		t.Fatalf("create %s: %v", name, err)
	}
	t.Cleanup(func() {
		if err := wginterface.Delete(name); err != nil {
			log.Printf("Failed to delete interface %s: %v", name, err)
		}
	})

	dev, err := wginterface.Open(name)
	if err != nil {
		t.Fatalf("open %s: %v", name, err)
	}
	t.Cleanup(func() { dev.Close() })

	if err := dev.Configure(privKey, port); err != nil {
		t.Fatalf("configure %s: %v", name, err)
	}

	return &testPeer{
		Name:    name,
		VpnIP:   vpnIP,
		Port:    port,
		PrivKey: privKey,
		PubKey:  privKey.PublicKey(),
		Dev:     dev,
	}
}

// waitHandshake polls until the named peer has completed a handshake or the timeout elapses.
func waitHandshake(t *testing.T, dev *wginterface.Device, pubKey wgtypes.Key, timeout time.Duration) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		p, err := dev.Peer(pubKey)
		if err != nil {
			t.Fatalf("peer lookup: %v", err)
		}
		if !p.LastHandshakeTime.IsZero() {
			return
		}
		time.Sleep(200 * time.Millisecond)
	}
	t.Fatalf("no handshake within %v", timeout)
}

func TestDirectHandshake(t *testing.T) {
	p1 := newTestPeer(t, "wgtest0", netip.MustParseAddr("192.168.99.1"), testBasePort)
	p2 := newTestPeer(t, "wgtest1", netip.MustParseAddr("192.168.99.2"), testBasePort+1)

	if err := p1.Dev.AddDirect(p2.PubKey, p2.Endpoint(), p2.VpnIP); err != nil {
		t.Fatalf("p1 AddDirect: %v", err)
	}
	if err := p2.Dev.AddDirect(p1.PubKey, p1.Endpoint(), p1.VpnIP); err != nil {
		t.Fatalf("p2 AddDirect: %v", err)
	}

	waitHandshake(t, p1.Dev, p2.PubKey, 30*time.Second)
	waitHandshake(t, p2.Dev, p1.PubKey, 30*time.Second)
}

func TestProbeAndPromote(t *testing.T) {
	p1 := newTestPeer(t, "wgtest0", netip.MustParseAddr("192.168.99.1"), testBasePort)
	p2 := newTestPeer(t, "wgtest1", netip.MustParseAddr("192.168.99.2"), testBasePort+1)

	// p2 needs a peer entry for p1 so it can respond to the handshake initiation.
	if err := p2.Dev.AddDirect(p1.PubKey, p1.Endpoint(), p1.VpnIP); err != nil {
		t.Fatalf("p2 AddDirect: %v", err)
	}

	if err := p1.Dev.AddProbe(p2.PubKey, p2.Endpoint()); err != nil {
		t.Fatalf("AddProbe: %v", err)
	}
	waitHandshake(t, p1.Dev, p2.PubKey, 30*time.Second)

	if err := p1.Dev.Promote(p2.PubKey, p2.VpnIP); err != nil {
		t.Fatalf("Promote: %v", err)
	}

	peer, err := p1.Dev.Peer(p2.PubKey)
	if err != nil {
		t.Fatalf("Peer: %v", err)
	}
	checkAllowedIP(t, peer, p2.VpnIP, 32)
}

func TestRelayHandshakes(t *testing.T) {
	vpnNetwork := netip.MustParsePrefix("192.168.99.0/24")

	relay := newTestPeer(t, "wgtest0", netip.MustParseAddr("192.168.99.1"), testBasePort)
	peer1 := newTestPeer(t, "wgtest1", netip.MustParseAddr("192.168.99.2"), testBasePort+1)
	peer2 := newTestPeer(t, "wgtest2", netip.MustParseAddr("192.168.99.3"), testBasePort+2)

	if err := relay.Dev.AddDirect(peer1.PubKey, peer1.Endpoint(), peer1.VpnIP); err != nil {
		t.Fatalf("relay AddDirect peer1: %v", err)
	}
	if err := relay.Dev.AddDirect(peer2.PubKey, peer2.Endpoint(), peer2.VpnIP); err != nil {
		t.Fatalf("relay AddDirect peer2: %v", err)
	}
	if err := peer1.Dev.SetRelay(relay.PubKey, relay.Endpoint(), vpnNetwork); err != nil {
		t.Fatalf("peer1 SetRelay: %v", err)
	}
	if err := peer2.Dev.SetRelay(relay.PubKey, relay.Endpoint(), vpnNetwork); err != nil {
		t.Fatalf("peer2 SetRelay: %v", err)
	}

	waitHandshake(t, relay.Dev, peer1.PubKey, 30*time.Second)
	waitHandshake(t, relay.Dev, peer2.PubKey, 30*time.Second)
	waitHandshake(t, peer1.Dev, relay.PubKey, 30*time.Second)
	waitHandshake(t, peer2.Dev, relay.PubKey, 30*time.Second)

	// relay has /32 entries for each peer
	p, err := relay.Dev.Peer(peer1.PubKey)
	if err != nil {
		t.Fatalf("relay peer1: %v", err)
	}
	checkAllowedIP(t, p, peer1.VpnIP, 32)

	p, err = relay.Dev.Peer(peer2.PubKey)
	if err != nil {
		t.Fatalf("relay peer2: %v", err)
	}
	checkAllowedIP(t, p, peer2.VpnIP, 32)

	// peers have /24 fallback route via relay
	p, err = peer1.Dev.Peer(relay.PubKey)
	if err != nil {
		t.Fatalf("peer1 relay: %v", err)
	}
	checkAllowedIP(t, p, vpnNetwork.Masked().Addr(), 24)

	p, err = peer2.Dev.Peer(relay.PubKey)
	if err != nil {
		t.Fatalf("peer2 relay: %v", err)
	}
	checkAllowedIP(t, p, vpnNetwork.Masked().Addr(), 24)
}

func TestRemovePeer(t *testing.T) {
	p1 := newTestPeer(t, "wgtest0", netip.MustParseAddr("192.168.99.1"), testBasePort)
	p2 := newTestPeer(t, "wgtest1", netip.MustParseAddr("192.168.99.2"), testBasePort+1)

	if err := p1.Dev.AddDirect(p2.PubKey, p2.Endpoint(), p2.VpnIP); err != nil {
		t.Fatalf("AddDirect: %v", err)
	}
	if err := p2.Dev.AddDirect(p1.PubKey, p1.Endpoint(), p1.VpnIP); err != nil {
		t.Fatalf("AddDirect: %v", err)
	}
	waitHandshake(t, p1.Dev, p2.PubKey, 30*time.Second)

	if err := p1.Dev.RemovePeer(p2.PubKey); err != nil {
		t.Fatalf("RemovePeer: %v", err)
	}
	if _, err := p1.Dev.Peer(p2.PubKey); err == nil {
		t.Fatal("expected error after RemovePeer, got nil")
	}
}

func TestEnableForwarding(t *testing.T) {
	p := newTestPeer(t, "wgtest0", netip.MustParseAddr("192.168.99.1"), testBasePort)

	if err := p.Dev.EnableForwarding(); err != nil {
		t.Fatalf("EnableForwarding: %v", err)
	}

	data, err := os.ReadFile(fmt.Sprintf("/proc/sys/net/ipv4/conf/%s/forwarding", p.Name))
	if err != nil {
		t.Fatalf("read forwarding: %v", err)
	}
	if strings.TrimSpace(string(data)) != "1" {
		t.Fatalf("expected forwarding=1, got %q", string(data))
	}
}

func TestPromoteKeepalive(t *testing.T) {
	p1 := newTestPeer(t, "wgtest0", netip.MustParseAddr("192.168.99.1"), testBasePort)
	p2 := newTestPeer(t, "wgtest1", netip.MustParseAddr("192.168.99.2"), testBasePort+1)

	if err := p2.Dev.AddDirect(p1.PubKey, p1.Endpoint(), p1.VpnIP); err != nil {
		t.Fatalf("p2 AddDirect: %v", err)
	}
	if err := p1.Dev.AddProbe(p2.PubKey, p2.Endpoint()); err != nil {
		t.Fatalf("AddProbe: %v", err)
	}
	waitHandshake(t, p1.Dev, p2.PubKey, 30*time.Second)

	if err := p1.Dev.Promote(p2.PubKey, p2.VpnIP); err != nil {
		t.Fatalf("Promote: %v", err)
	}

	peer, err := p1.Dev.Peer(p2.PubKey)
	if err != nil {
		t.Fatalf("Peer: %v", err)
	}
	if peer.PersistentKeepaliveInterval != 0 {
		t.Fatalf("expected keepalive disabled after promote, got %v", peer.PersistentKeepaliveInterval)
	}
}

func TestPeersCount(t *testing.T) {
	relay := newTestPeer(t, "wgtest0", netip.MustParseAddr("192.168.99.1"), testBasePort)
	peer1 := newTestPeer(t, "wgtest1", netip.MustParseAddr("192.168.99.2"), testBasePort+1)
	peer2 := newTestPeer(t, "wgtest2", netip.MustParseAddr("192.168.99.3"), testBasePort+2)

	if err := relay.Dev.AddDirect(peer1.PubKey, peer1.Endpoint(), peer1.VpnIP); err != nil {
		t.Fatalf("AddDirect peer1: %v", err)
	}
	if err := relay.Dev.AddDirect(peer2.PubKey, peer2.Endpoint(), peer2.VpnIP); err != nil {
		t.Fatalf("AddDirect peer2: %v", err)
	}

	peers, err := relay.Dev.Peers()
	if err != nil {
		t.Fatalf("Peers: %v", err)
	}
	if len(peers) != 2 {
		t.Fatalf("expected 2 peers, got %d", len(peers))
	}

	if err := relay.Dev.RemovePeer(peer1.PubKey); err != nil {
		t.Fatalf("RemovePeer: %v", err)
	}

	peers, err = relay.Dev.Peers()
	if err != nil {
		t.Fatalf("Peers after remove: %v", err)
	}
	if len(peers) != 1 {
		t.Fatalf("expected 1 peer after remove, got %d", len(peers))
	}
	if peers[0].PublicKey != peer2.PubKey {
		t.Fatal("wrong peer remained after remove")
	}
}

// checkAllowedIP asserts that a peer has exactly one AllowedIP matching addr/bits.
func checkAllowedIP(t *testing.T, p wgtypes.Peer, addr netip.Addr, bits int) {
	t.Helper()
	if len(p.AllowedIPs) != 1 {
		t.Fatalf("expected 1 AllowedIP, got %d", len(p.AllowedIPs))
	}
	ones, _ := p.AllowedIPs[0].Mask.Size()
	if ones != bits {
		t.Fatalf("expected /%d, got /%d", bits, ones)
	}
	got := netip.AddrFrom4([4]byte(p.AllowedIPs[0].IP.To4()))
	if got != addr {
		t.Fatalf("expected AllowedIP %v, got %v", addr, got)
	}
}