package nftables

import (
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"strings"

	"github.com/wild-cloud/wild-central/internal/storage"
)

const defaultRulesPath = "/etc/nftables.d/wild-cloud.nft"

// Manager handles nftables rule generation and application
type Manager struct {
	rulesPath string
}

// NewManager creates a new nftables manager
func NewManager(rulesPath string) *Manager {
	if rulesPath == "" {
		rulesPath = defaultRulesPath
	}
	return &Manager{rulesPath: rulesPath}
}

// GetRulesPath returns the nftables rules file path
func (m *Manager) GetRulesPath() string {
	return m.rulesPath
}

// Generate creates nftables rules that track HAProxy listening ports plus any extra ports.
// Policy is always "accept" — the allowed_tcp_ports set documents what is in use
// and enables strict WAN filtering when wanInterface is configured.
// extraTCPPorts and extraUDPPorts are user-specified ports to allow beyond the HAProxy set.
func (m *Manager) Generate(haproxyPorts []int, extraTCPPorts []int, extraUDPPorts []int, wanInterface string) string {
	// Merge and deduplicate TCP ports (HAProxy + extra TCP)
	seen := make(map[int]bool)
	var tcpPorts []int
	for _, p := range haproxyPorts {
		if !seen[p] {
			seen[p] = true
			tcpPorts = append(tcpPorts, p)
		}
	}
	for _, p := range extraTCPPorts {
		if !seen[p] {
			seen[p] = true
			tcpPorts = append(tcpPorts, p)
		}
	}

	// Build UDP exception list: DNS + DHCP hardcoded, plus user extras
	udpBase := []int{53, 67, 68}
	seenUDP := make(map[int]bool)
	for _, p := range udpBase {
		seenUDP[p] = true
	}
	udpPorts := append([]int{}, udpBase...)
	for _, p := range extraUDPPorts {
		if !seenUDP[p] {
			seenUDP[p] = true
			udpPorts = append(udpPorts, p)
		}
	}

	var sb strings.Builder

	sb.WriteString("# Wild Cloud nftables rules\n")
	sb.WriteString("# Managed by Wild Cloud Central API — do not edit manually\n\n")

	// Delete and recreate the table so re-applying never accumulates stale set elements.
	// The first line ensures the table exists (so delete never errors on first run).
	sb.WriteString("table inet wild-cloud {}\n")
	sb.WriteString("delete table inet wild-cloud\n\n")

	sb.WriteString("table inet wild-cloud {\n")
	sb.WriteString("    # TCP ports allowed through the firewall (HAProxy + extra TCP)\n")
	sb.WriteString("    set allowed_tcp_ports {\n")
	sb.WriteString("        type inet_service\n")
	sb.WriteString("        flags interval\n")
	if len(tcpPorts) > 0 {
		fmt.Fprintf(&sb, "        elements = { %s }\n", portsToStr(tcpPorts))
	}
	sb.WriteString("    }\n\n")

	sb.WriteString("    # UDP ports allowed through the firewall (DNS/DHCP + extra UDP)\n")
	sb.WriteString("    set allowed_udp_ports {\n")
	sb.WriteString("        type inet_service\n")
	sb.WriteString("        flags interval\n")
	if len(udpPorts) > 0 {
		fmt.Fprintf(&sb, "        elements = { %s }\n", portsToStr(udpPorts))
	}
	sb.WriteString("    }\n\n")

	sb.WriteString("    chain input {\n")
	sb.WriteString("        type filter hook input priority filter; policy accept;\n\n")
	sb.WriteString("        iif \"lo\" accept\n")
	sb.WriteString("        ct state { established, related } accept\n")

	if wanInterface != "" {
		sb.WriteString("\n        # Strict WAN filtering: only allow listed TCP and UDP ports\n")
		fmt.Fprintf(&sb, "        iif \"%s\" tcp dport != @allowed_tcp_ports drop\n", wanInterface)
		fmt.Fprintf(&sb, "        iif \"%s\" udp dport != @allowed_udp_ports drop\n", wanInterface)
	}

	sb.WriteString("    }\n")
	sb.WriteString("}\n")

	return sb.String()
}

// portsToStr converts a slice of ports to a comma-separated string
func portsToStr(ports []int) string {
	strs := make([]string, len(ports))
	for i, p := range ports {
		strs[i] = fmt.Sprintf("%d", p)
	}
	return strings.Join(strs, ", ")
}

// ValidateRules tests an nftables rules file for syntax errors.
// Uses sudo to allow the wildcloud user to run the read-only check.
func (m *Manager) ValidateRules(rulesPath string) error {
	cmd := exec.Command("sudo", "nft", "-c", "-f", rulesPath)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("nftables validation failed: %w (output: %s)", err, string(output))
	}
	return nil
}

// SafeApply validates, backs up, writes, and applies nftables rules.
// On apply failure, rolls back to the previous rules.
func (m *Manager) SafeApply(content string) error {
	// Validate syntax
	tmpFile := "/tmp/wild-cloud-nft-safeapply.tmp"
	if err := os.WriteFile(tmpFile, []byte(content), 0644); err != nil {
		return fmt.Errorf("writing validation file: %w", err)
	}
	defer os.Remove(tmpFile)

	if err := m.ValidateRules(tmpFile); err != nil {
		return fmt.Errorf("validation failed: %w", err)
	}

	// Backup current rules
	if storage.FileExists(m.rulesPath) {
		_ = storage.CopyFile(m.rulesPath, m.rulesPath+".bak")
	}

	// Write atomically
	if err := storage.WriteFileAtomic(m.rulesPath, []byte(content), 0644); err != nil {
		return fmt.Errorf("write failed: %w", err)
	}

	// Apply to kernel
	if err := m.ApplyRules(); err != nil {
		m.rollback()
		return fmt.Errorf("apply failed (rolled back): %w", err)
	}

	// Verify rules loaded
	if err := m.verify(); err != nil {
		m.rollback()
		return fmt.Errorf("verification failed (rolled back): %w", err)
	}

	return nil
}

func (m *Manager) rollback() {
	bakPath := m.rulesPath + ".bak"
	if storage.FileExists(bakPath) {
		_ = os.Rename(bakPath, m.rulesPath)
		_ = m.ApplyRules()
		slog.Warn("rolled back to previous rules", "component", "nftables", "path", m.rulesPath)
	}
}

func (m *Manager) verify() error {
	status, err := m.GetStatus()
	if err != nil || status == "" {
		return fmt.Errorf("nftables table not loaded after apply")
	}
	return nil
}

// WriteRules validates the content then writes the nftables rules file.
// Validation uses a temp file in /tmp (world-writable) to avoid needing
// write permission on the /etc/nftables.d/ directory itself.
func (m *Manager) WriteRules(content string) error {
	tempFile := "/tmp/wild-cloud-nft-validate.tmp"

	if err := os.WriteFile(tempFile, []byte(content), 0644); err != nil {
		return fmt.Errorf("writing temp rules: %w", err)
	}
	defer os.Remove(tempFile)

	if err := m.ValidateRules(tempFile); err != nil {
		return err
	}

	if err := storage.WriteFileAtomic(m.rulesPath, []byte(content), 0644); err != nil {
		return fmt.Errorf("writing rules file: %w", err)
	}

	slog.Info("nftables rules written", "component", "nftables", "path", m.rulesPath)
	return nil
}

// ApplyRules loads the rules file into the kernel via a systemd oneshot service.
// The wildcloud user has polkit permission to start wild-cloud-nftables-reload.service.
func (m *Manager) ApplyRules() error {
	cmd := exec.Command("systemctl", "start", "wild-cloud-nftables-reload.service")
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("applying nftables rules: %w (output: %s)", err, string(output))
	}
	slog.Info("nftables rules applied", "component", "nftables")
	return nil
}

// WriteDisabledRules writes a rules file that flushes the wild-cloud table,
// removing all Wild Cloud firewall rules from the kernel when applied.
func (m *Manager) WriteDisabledRules() error {
	content := "# Wild Cloud nftables rules — firewall disabled\n" +
		"# Managed by Wild Cloud Central API — do not edit manually\n\n" +
		"table inet wild-cloud {}\n" +
		"delete table inet wild-cloud\n"
	if err := storage.WriteFileAtomic(m.rulesPath, []byte(content), 0644); err != nil {
		return fmt.Errorf("writing disabled rules file: %w", err)
	}
	slog.Info("nftables rules disabled", "component", "nftables", "path", m.rulesPath)
	return nil
}

// GetStatus returns the current wild-cloud nftables table as a string.
// Uses sudo to allow the wildcloud user to read kernel state.
func (m *Manager) GetStatus() (string, error) {
	cmd := exec.Command("sudo", "nft", "list", "table", "inet", "wild-cloud")
	output, err := cmd.CombinedOutput()
	if err != nil {
		// Table may not exist yet — not an error
		return "", nil
	}
	return string(output), nil
}