// Package firewall builds and (optionally) applies an nftables ruleset for
// a WireGuard server. Every server gets its own table
// (inet wireguard_ui_<serverID>) so applying/removing one server's rules
// never touches any other table on the system.
package firewall

import (
	"fmt"
	"strings"

	"github.com/ngoduykhanh/wireguard-ui/model"
)

// TableName returns the dedicated nftables table name for a server.
func TableName(serverID string) string {
	return "wireguard_ui_" + serverID
}

// buildRuleLine renders one custom rule as an nftables statement.
func buildRuleLine(rule model.FirewallRule) string {
	var parts []string
	if rule.Source != "" {
		parts = append(parts, fmt.Sprintf("ip saddr %s", rule.Source))
	}
	switch {
	case rule.Protocol != "" && rule.Port != "":
		parts = append(parts, fmt.Sprintf("%s dport %s", rule.Protocol, rule.Port))
	case rule.Protocol != "":
		parts = append(parts, fmt.Sprintf("meta l4proto %s", rule.Protocol))
	case rule.Port != "":
		parts = append(parts, fmt.Sprintf("th dport %s", rule.Port))
	}
	parts = append(parts, rule.Action)
	comment := rule.Comment
	if comment == "" {
		comment = "wg-ui-multi custom rule"
	}
	parts = append(parts, fmt.Sprintf("comment %q", comment))
	return "        " + strings.Join(parts, " ")
}

// GenerateRuleset renders the full nftables ruleset for a server: the
// baseline (listen-port accept, WireGuard-interface forwarding, optional
// LAN forwarding) plus every enabled custom rule, grouped by chain.
func GenerateRuleset(server model.Server, settings model.ServerSetting, rules []model.FirewallRule) string {
	ifaceName := "wgX"
	listenPort := 0
	if server.Interface != nil {
		if server.Interface.Name != "" {
			ifaceName = server.Interface.Name
		}
		listenPort = server.Interface.ListenPort
	}

	var inputExtra, forwardExtra []string
	for _, rule := range rules {
		if !rule.Enabled {
			continue
		}
		line := buildRuleLine(rule)
		if rule.Direction == "forward" {
			forwardExtra = append(forwardExtra, line)
		} else {
			inputExtra = append(inputExtra, line)
		}
	}

	var b strings.Builder
	fmt.Fprintf(&b, "# nftables ruleset for server %q (%s)\n", server.Name, server.ID)
	fmt.Fprintf(&b, "# Generated by wireguard-ui-multi.\n\n")

	fmt.Fprintf(&b, "table inet %s {\n", TableName(server.ID))
	fmt.Fprintf(&b, "    chain input {\n")
	fmt.Fprintf(&b, "        type filter hook input priority 0; policy accept;\n")
	fmt.Fprintf(&b, "        udp dport %d accept comment \"wg-ui-multi: %s\"\n", listenPort, server.ID)
	for _, line := range inputExtra {
		fmt.Fprintf(&b, "%s\n", line)
	}
	fmt.Fprintf(&b, "    }\n\n")

	fmt.Fprintf(&b, "    chain forward {\n")
	fmt.Fprintf(&b, "        type filter hook forward priority 0; policy accept;\n")
	fmt.Fprintf(&b, "        iifname \"%s\" accept comment \"wg-ui-multi: %s inbound\"\n", ifaceName, server.ID)
	fmt.Fprintf(&b, "        oifname \"%s\" accept comment \"wg-ui-multi: %s outbound\"\n", ifaceName, server.ID)
	if settings.LanInterface != "" {
		fmt.Fprintf(&b, "        iifname \"%s\" oifname \"%s\" accept comment \"wg-ui-multi: %s -> lan\"\n", ifaceName, settings.LanInterface, server.ID)
		fmt.Fprintf(&b, "        iifname \"%s\" oifname \"%s\" accept comment \"wg-ui-multi: lan -> %s\"\n", settings.LanInterface, ifaceName, server.ID)
	}
	for _, line := range forwardExtra {
		fmt.Fprintf(&b, "%s\n", line)
	}
	fmt.Fprintf(&b, "    }\n")
	fmt.Fprintf(&b, "}\n")

	return b.String()
}