package e2e

import (
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func TestDataPlaneConnectivity(t *testing.T) {
	// 1. Determine binary path
	binaryPath := EnsureBinary(t)

	// 2. Setup isolated config & runtime folders for testing
	tmpDir := t.TempDir()
	profile := "e2e-dataplane-test"

	// Create a dummy peer UDP listener inside our test harness
	// to simulate the remote WireGuard peer. We'll listen on a random port.
	addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("Failed to resolve UDP address: %v", err)
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		t.Fatalf("Failed to start mock remote WG UDP listener: %v", err)
	}
	defer func() { _ = conn.Close() }()

	localPort := conn.LocalAddr().(*net.UDPAddr).Port

	// Generate profile with valid Base64 keys
	clientPrivKey := "YXNkZmFzZGZhc2RmYXNkZmFzZGZhc2RmYXNkZmFzZGY=" // 32-bytes base64
	peerPubKey := "YXNkZmFzZGZhc2RmYXNkZmFzZGZhc2RmYXNkZmFzZGY="

	confContent := fmt.Sprintf(`[Interface]
PrivateKey = %s
Address = 10.0.0.2/24

[Peer]
PublicKey = %s
Endpoint = 127.0.0.1:%d
AllowedIPs = 10.0.0.0/24
`, clientPrivKey, peerPubKey, localPort)

	profilesDir := filepath.Join(tmpDir, "wg-wrap", "profiles")
	if err := os.MkdirAll(profilesDir, 0755); err != nil {
		t.Fatalf("Failed to create temporary profiles dir: %v", err)
	}
	profilePath := filepath.Join(profilesDir, profile+".conf")
	if err := os.WriteFile(profilePath, []byte(confContent), 0644); err != nil {
		t.Fatalf("Failed to write temporary test profile: %v", err)
	}

	// 3. Launch wg-wrap with a command that triggers traffic
	cmd := exec.Command(binaryPath, "--profile", profile, "--", "ping", "-c", "1", "-W", "1", "10.0.0.1")
	cmd.Env = append(os.Environ(),
		fmt.Sprintf("XDG_CONFIG_HOME=%s", tmpDir),
		fmt.Sprintf("XDG_RUNTIME_DIR=%s", tmpDir),
	)

	packetChan := make(chan []byte, 1)
	go func() {
		buf := make([]byte, 2048)
		_ = conn.SetReadDeadline(time.Now().Add(3 * time.Second))
		n, _, err := conn.ReadFrom(buf)
		if err == nil && n > 0 {
			packetChan <- buf[:n]
		} else {
			packetChan <- nil
		}
	}()

	if err := cmd.Run(); err != nil {
		t.Logf("wg-wrap command returned error (expected since mock peer doesn't reply): %v", err)
	}

	select {
	case packet := <-packetChan:
		if packet == nil {
			t.Error("Mock remote WG UDP listener did not receive any packet from wg-wrap")
		} else {
			t.Logf("Mock remote WG UDP listener successfully received packet of size %d", len(packet))
		}
	case <-time.After(4 * time.Second):
		t.Error("Timed out waiting for WireGuard packet from wg-wrap")
	}

	t.Log("Successfully created tunnel namespace and ran isolated command rootlessly.")
}

func TestNetworkIsolation(t *testing.T) {
	binaryPath := EnsureBinary(t)

	cmd := exec.Command(binaryPath, "test-ns")
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("wg-wrap test-ns failed: %v\nOutput: %s", err, string(out))
	}

	if !strings.Contains(string(out), "Isolation Verified: OK") {
		t.Errorf("Expected 'Isolation Verified: OK', got: %q", string(out))
	}
}

func TestDNSIsolation(t *testing.T) {
	binaryPath := EnsureBinary(t)

	// 1. Start Mock DNS Server
	dnsServer, port := StartMockDNSServer(t)
	defer dnsServer.Close()

	// 2. Setup isolated config
	tmpDir := t.TempDir()
	profile := "test-dns-isolation"
	clientPrivKey := "YXNkZmFzZGZhc2RmYXNkZmFzZGZhc2RmYXNkZmFzZGY="
	peerPubKey := "YXNkZmFzZGZhc2RmYXNkZmFzZGZhc2RmYXNkZmFzZGY="
	dnsServerIP := "10.0.0.1"

	confContent := fmt.Sprintf(`[Interface]
PrivateKey = %s
Address = 10.0.0.2/24

[Peer]
PublicKey = %s
Endpoint = 127.0.0.1:%d
AllowedIPs = 10.0.0.0/24
`, clientPrivKey, peerPubKey, port)

	profilesDir := filepath.Join(tmpDir, "wg-wrap", "profiles")
	_ = os.MkdirAll(profilesDir, 0755)
	profilePath := filepath.Join(profilesDir, profile+".conf")
	_ = os.WriteFile(profilePath, []byte(confContent), 0644)

	// 3. Test /etc/resolv.conf modification
	expectedDNS := "1.1.1.1"
	cmd := exec.Command(binaryPath, "--profile", profile, "--dns-server", expectedDNS, "--", "cat", "/etc/resolv.conf")
	cmd.Env = append(os.Environ(),
		fmt.Sprintf("XDG_CONFIG_HOME=%s", tmpDir),
		fmt.Sprintf("XDG_RUNTIME_DIR=%s", tmpDir),
	)

	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("Failed to run resolv.conf check: %v\nOutput: %s", err, string(out))
	}

	if !strings.Contains(string(out), "nameserver "+expectedDNS) {
		t.Errorf("Expected /etc/resolv.conf to contain %s, but got: %q", expectedDNS, string(out))
	}

	// 4. Test Data Path: Send a ping to trigger Handshake on the mock server
	cmdQuery := exec.Command(binaryPath, "--profile", profile, "--", "ping", "-c", "1", "-W", "1", dnsServerIP)
	cmdQuery.Env = cmd.Env

	packetReceived := make(chan bool, 1)
	go func() {
		success, _ := dnsServer.ListenAndRespond(5 * time.Second)
		packetReceived <- <-success
	}()

	if err := cmdQuery.Run(); err != nil {
		t.Logf("Note: query command failed as expected (since we didn't implement a full DNS stack), but we check if packet arrived: %v", err)
	}

	select {
	case received := <-packetReceived:
		if !received {
			t.Error("Mock DNS server did not receive the UDP packet through the tunnel")
		}
	case <-time.After(5 * time.Second):
		t.Error("Timed out waiting for DNS packet to reach mock server")
	}
}

func TestDNSPrecedence(t *testing.T) {
	binaryPath := EnsureBinary(t)

	tmpDir := t.TempDir()
	profileName := "test-dns-precedence"
	clientPrivKey := "YXNkZmFzZGZhc2RmYXNkZmFzZGZhc2RmYXNkZmFzZGY="
	peerPubKey := "YXNkZmFzZGZhc2RmYXNkZmFzZGZhc2RmYXNkZmFzZGY="

	tests := []struct {
		name        string
		configDNS   string
		cliDNS      string
		expectedDNS string
	}{
		{
			name:        "Fallback to safe DNS (1.1.1.1) when none is specified",
			configDNS:   "",
			cliDNS:      "",
			expectedDNS: "1.1.1.1",
		},
		{
			name:        "Use .conf specified DNS when no CLI flag is provided",
			configDNS:   "8.8.4.4",
			cliDNS:      "",
			expectedDNS: "8.8.4.4",
		},
		{
			name:        "CLI flag overrides .conf specified DNS",
			configDNS:   "8.8.4.4",
			cliDNS:      "9.9.9.9",
			expectedDNS: "9.9.9.9",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Write the profile conf with or without the DNS field
			var dnsLine string
			if tt.configDNS != "" {
				dnsLine = "DNS = " + tt.configDNS
			}

			confContent := fmt.Sprintf(`[Interface]
PrivateKey = %s
Address = 10.0.0.2/24
%s

[Peer]
PublicKey = %s
Endpoint = 127.0.0.1:51820
AllowedIPs = 10.0.0.0/24
`, clientPrivKey, dnsLine, peerPubKey)

			profilesDir := filepath.Join(tmpDir, "wg-wrap", "profiles")
			_ = os.MkdirAll(profilesDir, 0755)
			profilePath := filepath.Join(profilesDir, profileName+".conf")
			_ = os.WriteFile(profilePath, []byte(confContent), 0644)

			// Prepare command args
			args := []string{"--profile", profileName}
			if tt.cliDNS != "" {
				args = append(args, "--dns-server", tt.cliDNS)
			}
			args = append(args, "--", "cat", "/etc/resolv.conf")

			cmd := exec.Command(binaryPath, args...)
			cmd.Env = append(os.Environ(),
				fmt.Sprintf("XDG_CONFIG_HOME=%s", tmpDir),
				fmt.Sprintf("XDG_RUNTIME_DIR=%s", tmpDir),
			)

			out, err := cmd.CombinedOutput()
			if err != nil {
				t.Fatalf("Failed to execute resolv.conf check: %v\nOutput: %s", err, string(out))
			}

			if !strings.Contains(string(out), "nameserver "+tt.expectedDNS) {
				t.Errorf("Expected /etc/resolv.conf to contain nameserver %s, but got: %q", tt.expectedDNS, string(out))
			}
		})
	}
}

func TestMTUFragmentation(t *testing.T) {
	binaryPath := EnsureBinary(t)

	cmd := exec.Command(binaryPath, "--profile", "default", "--", "true")
	if err := cmd.Run(); err != nil {
		t.Errorf("expected command to pass, got: %v", err)
	}
}

func TestExitCodePropagation(t *testing.T) {
	binaryPath := EnsureBinary(t)

	// Run a command that exits with code 42
	cmd := exec.Command(binaryPath, "--profile", "default", "--", "sh", "-c", "exit 42")
	err := cmd.Run()
	if err == nil {
		t.Fatalf("expected command to fail with exit status 42, but it succeeded")
	}

	exitErr, ok := err.(*exec.ExitError)
	if !ok {
		t.Fatalf("expected error of type *exec.ExitError, got %T: %v", err, err)
	}

	if exitErr.ExitCode() != 42 {
		t.Errorf("expected exit code 42, got %d", exitErr.ExitCode())
	}
}