package firewall

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
)

const tableName = "nano-router"

// TproxyPort is the fixed Mihomo transparent-proxy port, always enabled.
const TproxyPort = 7893

// TproxyMark is the fwmark set on packets that should be routed to tproxy.
const TproxyMark = 1

// TproxyTable is the ip routing table used for fwmark-based local delivery.
const TproxyTable = 100

// Rule is a single stateless forward-filter rule.
type Rule struct {
	ID       string `yaml:"id"        json:"id"`
	Enabled  bool   `yaml:"enabled"   json:"enabled"`
	Action   string `yaml:"action"    json:"action"`    // accept | drop | reject
	Protocol string `yaml:"protocol"  json:"protocol"`  // tcp | udp | icmp | all
	SrcAddr  string `yaml:"src_addr"  json:"src_addr"`  // CIDR or IP, empty = any
	SrcPort  string `yaml:"src_port"  json:"src_port"`  // "80" | "80-443", empty = any
	DstAddr  string `yaml:"dst_addr"  json:"dst_addr"`
	DstPort  string `yaml:"dst_port"  json:"dst_port"`
	InIface  string `yaml:"in_iface"  json:"in_iface"`  // input interface, empty = any
	OutIface string `yaml:"out_iface" json:"out_iface"` // output interface, empty = any
	Comment  string `yaml:"comment"   json:"comment"`
}

// Config is the top-level firewall config stored in config.yaml.
type Config struct {
	Rules         []Rule `yaml:"rules"          json:"rules"`
	VLANIsolation bool   `yaml:"vlan_isolation" json:"vlan_isolation"`
}

// NATConfig holds NAT masquerade settings (passed to avoid a direct nat import).
type NATConfig struct {
	Interfaces []string
}

// ClientPolicies holds per-policy IP lists derived from known device policies.
type ClientPolicies struct {
	DisabledIPs []string // drop all traffic for these IPs
	VPNIPs      []string // redirect TCP/UDP to tproxy for these IPs
}

// IsInstalled reports whether the nft binary is available.
func IsInstalled() bool {
	_, err := exec.LookPath("nft")
	return err == nil
}

// CleanupAll removes all kernel state managed by nano-router so that a fresh
// applyConfig() starts from a guaranteed clean slate:
//   - All nftables tables created by nano-router (main + traffic)
//   - The ip rule and route used for tproxy fwmark delivery
//
// Errors are silently ignored — most arise because entries don't exist yet.
func CleanupAll() {
	// Remove nftables tables. The traffic table is managed by the traffic tracker.
	exec.Command("nft", "delete", "table", "ip", tableName).Run()
	exec.Command("nft", "delete", "table", "ip", "nano-router-nat").Run()
	exec.Command("nft", "delete", "table", "ip", "nano-router-traffic").Run()

	// Remove all ip rules pointing to our tproxy routing table.
	// Loop to handle duplicate entries that can accumulate across restarts.
	mark := fmt.Sprintf("%d", TproxyMark)
	table := fmt.Sprintf("%d", TproxyTable)
	for {
		err := exec.Command("ip", "rule", "del", "fwmark", mark, "table", table).Run()
		if err != nil {
			break
		}
	}

	// Remove the local default route in our tproxy routing table.
	exec.Command("ip", "route", "del",
		"local", "0.0.0.0/0", "dev", "lo", "table", table,
	).Run()
}

// SetupTproxyRouting installs the ip rule and route needed for TPROXY to work.
// Packets marked with TproxyMark are routed to the loopback interface (local delivery),
// which allows Mihomo's tproxy socket to receive them.
// Call CleanupAll first to avoid duplicate entries.
func SetupTproxyRouting() {
	mark := fmt.Sprintf("%d", TproxyMark)
	table := fmt.Sprintf("%d", TproxyTable)
	exec.Command("ip", "rule", "add", "fwmark", mark, "table", table).Run()
	exec.Command("ip", "route", "add",
		"local", "0.0.0.0/0", "dev", "lo", "table", table,
	).Run()
}

// ApplyAll atomically regenerates the complete nftables ruleset:
//   - Tproxy prerouting for cp.VPNIPs  (mangle priority, TPROXY → Mihomo :TproxyPort)
//   - NAT masquerade for natCfg.Interfaces
//   - Disabled client IP drops (cp.DisabledIPs)
//   - User rules from fwCfg (in order, enabled only)
//   - LAN isolation (if fwCfg.VLANIsolation): blocks traffic between any two LAN interfaces
//   - Default accept from LAN interfaces to WAN
//
// lanIfaces is the union of NAT interfaces and all VLAN interfaces.
func ApplyAll(natCfg NATConfig, fwCfg Config, lanIfaces []string, cp ClientPolicies) error {
	if err := os.WriteFile("/proc/sys/net/ipv4/ip_forward", []byte("1"), 0644); err != nil {
		return fmt.Errorf("enable ip_forward: %w", err)
	}

	// Remove both old and new table names to ensure clean state.
	exec.Command("nft", "delete", "table", "ip", "nano-router-nat").Run()
	exec.Command("nft", "delete", "table", "ip", tableName).Run()

	var activeRules []Rule
	for _, r := range fwCfg.Rules {
		if r.Enabled {
			activeRules = append(activeRules, r)
		}
	}

	hasNAT := len(natCfg.Interfaces) > 0
	hasDisabled := len(cp.DisabledIPs) > 0
	hasVPN := len(cp.VPNIPs) > 0
	hasVLANIsolation := fwCfg.VLANIsolation && len(lanIfaces) >= 2

	if !hasNAT && !hasDisabled && !hasVPN && !hasVLANIsolation && len(activeRules) == 0 {
		return nil
	}

	var sb strings.Builder
	fmt.Fprintf(&sb, "table ip %s {\n", tableName)

	// ── Tproxy prerouting chain (mangle priority) ────────────────────────────
	// Intercepts TCP/UDP from VPN-policy clients and delivers them to Mihomo.
	// Runs before the forward chain, so forwarding rules don't affect these packets.
	if hasVPN {
		sb.WriteString("    chain tproxy_pre {\n")
		sb.WriteString("        type filter hook prerouting priority mangle; policy accept;\n")
		// Never redirect traffic destined for the router itself (admin panel,
		// SSH, DNS served locally, etc.). Without this, Mihomo intercepts
		// connections to the router's own IPs and the admin UI becomes unreachable.
		sb.WriteString("        fib daddr type local return\n")
		for _, ip := range cp.VPNIPs {
			fmt.Fprintf(&sb,
				"        ip saddr %s meta l4proto { tcp, udp } tproxy to :%d meta mark set %d\n",
				ip, TproxyPort, TproxyMark,
			)
		}
		sb.WriteString("    }\n")
	}

	// ── Forward chain ────────────────────────────────────────────────────────
	sb.WriteString("    chain forward {\n")
	sb.WriteString("        type filter hook forward priority filter; policy drop;\n")
	sb.WriteString("        ct state established,related accept\n")

	// Drop traffic for disabled clients (both directions).
	for _, ip := range cp.DisabledIPs {
		fmt.Fprintf(&sb, "        ip saddr %s drop\n", ip)
		fmt.Fprintf(&sb, "        ip daddr %s drop\n", ip)
	}

	// User-defined forward rules (ordered, enabled only).
	for _, rule := range activeRules {
		line := buildRuleLine(rule)
		if line == "" {
			continue
		}
		if rule.Comment != "" {
			fmt.Fprintf(&sb, "        # %s\n", rule.Comment)
		}
		fmt.Fprintf(&sb, "        %s\n", line)
	}

	// LAN isolation — drop traffic between any two local (LAN) interfaces.
	// Placed AFTER user rules so explicit allow rules above still take effect.
	if hasVLANIsolation {
		quoted := make([]string, len(lanIfaces))
		for i, v := range lanIfaces {
			quoted[i] = fmt.Sprintf("%q", v)
		}
		set := "{ " + strings.Join(quoted, ", ") + " }"
		fmt.Fprintf(&sb, "        iifname %s oifname %s drop\n", set, set)
	}

	// Allow from LAN/VLAN interfaces outbound to WAN.
	for _, iface := range natCfg.Interfaces {
		fmt.Fprintf(&sb, "        iifname %q accept\n", iface)
	}

	sb.WriteString("    }\n")

	// ── Postrouting (masquerade) ─────────────────────────────────────────────
	if hasNAT {
		sb.WriteString("    chain postrouting {\n")
		sb.WriteString("        type nat hook postrouting priority srcnat; policy accept;\n")
		for _, iface := range natCfg.Interfaces {
			fmt.Fprintf(&sb, "        iifname %q masquerade\n", iface)
		}
		sb.WriteString("    }\n")
	}

	sb.WriteString("}\n")

	cmd := exec.Command("nft", "-f", "-")
	cmd.Stdin = strings.NewReader(sb.String())
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("nft apply: %s: %w", strings.TrimSpace(string(out)), err)
	}

	// Set up ip rule/route for tproxy fwmark routing when VPN clients are active.
	if hasVPN {
		SetupTproxyRouting()
	}

	// Flush connection tracking so existing sessions are re-evaluated.
	flushConntrack()

	return nil
}

// flushConntrack clears the kernel connection tracking table so that all traffic
// is re-evaluated against the current nftables ruleset.
func flushConntrack() {
	if err := exec.Command("conntrack", "-F").Run(); err == nil {
		return
	}
	_ = os.WriteFile("/proc/sys/net/netfilter/nf_conntrack_flush", []byte("1"), 0644)
}

// buildRuleLine converts a Rule to a single nftables match+action string.
// Returns "" if the rule has no valid action.
func buildRuleLine(r Rule) string {
	if r.Action == "" {
		return ""
	}

	var parts []string

	if r.InIface != "" {
		parts = append(parts, fmt.Sprintf("iifname %q", r.InIface))
	}
	if r.OutIface != "" {
		parts = append(parts, fmt.Sprintf("oifname %q", r.OutIface))
	}
	if r.SrcAddr != "" {
		parts = append(parts, "ip saddr "+r.SrcAddr)
	}
	if r.DstAddr != "" {
		parts = append(parts, "ip daddr "+r.DstAddr)
	}

	proto := strings.ToLower(r.Protocol)
	switch proto {
	case "tcp", "udp":
		if r.SrcPort != "" || r.DstPort != "" {
			if r.SrcPort != "" {
				parts = append(parts, proto+" sport "+r.SrcPort)
			}
			if r.DstPort != "" {
				parts = append(parts, proto+" dport "+r.DstPort)
			}
		} else {
			parts = append(parts, "ip protocol "+proto)
		}
	case "icmp":
		parts = append(parts, "ip protocol icmp")
	}

	parts = append(parts, r.Action)
	return strings.Join(parts, " ")
}