// Copyright 2026 Cisco Systems, Inc. and its affiliates
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

// Package iptables generates Linux iptables rules from a DefenseClaw firewall config.
// No privileged operations — compile only. Applying rules is the administrator's job.
package iptables

import (
	"fmt"
	"net"
	"regexp"
	"strings"

	"github.com/defenseclaw/defenseclaw/internal/firewall"
)

// Compiler generates iptables rules for Linux.
type Compiler struct{}

func New() *Compiler { return &Compiler{} }

func (c *Compiler) Platform() string { return "iptables" }

var safeArgPattern = regexp.MustCompile(`^[a-zA-Z0-9._:/%\-]+$`)

var knownFlags = map[string]bool{
	"-A": true, "-D": true, "-I": true, "-F": true, "-P": true, "-N": true, "-X": true,
	"-j": true, "-d": true, "-s": true, "-p": true, "-m": true,
	"--dport": true, "--sport": true, "--limit": true,
	"--log-prefix": true, "--log-level": true, "--state": true,
	"INPUT": true, "OUTPUT": true, "FORWARD": true,
	"ACCEPT": true, "DROP": true, "REJECT": true, "LOG": true,
	"tcp": true, "udp": true, "icmp": true,
	"limit": true, "state": true, "comment": true,
	"NEW": true, "ESTABLISHED": true, "RELATED": true, "INVALID": true,
}

func (c *Compiler) ValidateArg(arg string) error {
	if strings.ContainsAny(arg, ";|&`$(){}[]<>!~*?'\"\\") {
		return fmt.Errorf("iptables: argument contains shell metacharacters: %q", arg)
	}
	if knownFlags[arg] {
		return nil
	}
	if net.ParseIP(arg) != nil {
		return nil
	}
	if _, _, err := net.ParseCIDR(arg); err == nil {
		return nil
	}
	if matched, _ := regexp.MatchString(`^\d+/(sec|min|hour|day)$`, arg); matched {
		return nil
	}
	if matched, _ := regexp.MatchString(`^\d+$`, arg); matched {
		return nil
	}
	if matched, _ := regexp.MatchString(`^[\w\[\]\- ]+$`, arg); matched {
		return nil
	}
	if !safeArgPattern.MatchString(arg) {
		return fmt.Errorf("iptables: argument contains unsafe characters: %q", arg)
	}
	return nil
}

func (c *Compiler) ApplyCommand(rulesPath string) string {
	return fmt.Sprintf("sudo iptables-restore < %s", rulesPath)
}

func (c *Compiler) RemoveCommand() string {
	return "sudo iptables -F OUTPUT && sudo iptables -P OUTPUT ACCEPT"
}

// Compile generates iptables-restore format rules from the configuration.
func (c *Compiler) Compile(cfg *firewall.FirewallConfig) ([]string, error) {
	var rules []string

	rules = append(rules,
		"# DefenseClaw egress firewall rules — generated by defenseclaw firewall generate",
		fmt.Sprintf("# default_action: %s", cfg.DefaultAction),
		"*filter",
		":INPUT ACCEPT [0:0]",
		":FORWARD ACCEPT [0:0]",
		":OUTPUT ACCEPT [0:0]",
		"",
		"# Allow loopback",
		"-A OUTPUT -o lo -j ACCEPT",
		"",
		"# Allow DNS",
		"-A OUTPUT -p udp --dport 53 -j ACCEPT",
		"-A OUTPUT -p tcp --dport 53 -j ACCEPT",
		"",
		"# Allow established/related connections",
		"-A OUTPUT -m state --state ESTABLISHED,RELATED -j ACCEPT",
		"",
	)

	// Explicit named rules.
	for _, rule := range cfg.Rules {
		r, err := c.compileRule(rule)
		if err != nil {
			return nil, fmt.Errorf("iptables: compile rule %q: %w", rule.Name, err)
		}
		rules = append(rules, fmt.Sprintf("# Rule: %s", rule.Name), r)
	}

	// Allowlist: IPs.
	if len(cfg.Allowlist.IPs) > 0 {
		rules = append(rules, "", "# Allowlist: IPs")
		for _, ip := range cfg.Allowlist.IPs {
			if strings.Contains(ip, ":") {
				continue // Skip IPv6 for iptables (use ip6tables separately)
			}
			if len(cfg.Allowlist.Ports) == 0 {
				rules = append(rules, fmt.Sprintf("-A OUTPUT -d %s -j ACCEPT", ip))
				continue
			}
			for _, port := range cfg.Allowlist.Ports {
				rules = append(rules, fmt.Sprintf("-A OUTPUT -d %s -p tcp --dport %d -j ACCEPT", ip, port))
			}
		}
	}

	// Allowlist: domains (resolved at generation time).
	if len(cfg.Allowlist.Domains) > 0 {
		rules = append(rules, "", "# Allowlist: domains")
		for _, domain := range cfg.Allowlist.Domains {
			if strings.HasPrefix(domain, "*.") {
				rules = append(rules, fmt.Sprintf("# Wildcard domain (resolve manually): %s", domain))
				continue
			}
			ips, err := net.LookupHost(domain)
			if err != nil {
				rules = append(rules, fmt.Sprintf("# Failed to resolve: %s (%v)", domain, err))
				continue
			}
			for _, ip := range ips {
				if strings.Contains(ip, ":") {
					continue // Skip IPv6
				}
				if len(cfg.Allowlist.Ports) == 0 {
					rules = append(rules, fmt.Sprintf("-A OUTPUT -d %s -j ACCEPT # %s", ip, domain))
					continue
				}
				for _, port := range cfg.Allowlist.Ports {
					rules = append(rules, fmt.Sprintf("-A OUTPUT -d %s -p tcp --dport %d -j ACCEPT # %s", ip, port, domain))
				}
			}
		}
	}

	// Logging + default action.
	rules = append(rules, "")
	if cfg.Logging.Enabled && cfg.DefaultAction == "deny" {
		prefix := cfg.Logging.Prefix
		if len(prefix) > 29 {
			prefix = prefix[:29]
		}
		rules = append(rules, fmt.Sprintf(
			"-A OUTPUT -m limit --limit %s -j LOG --log-prefix \"%s \" --log-level 4",
			cfg.Logging.RateLimit, prefix,
		))
	}

	if cfg.DefaultAction == "deny" {
		rules = append(rules, "# Default: deny", "-A OUTPUT -j DROP")
	}

	rules = append(rules, "COMMIT")
	return rules, nil
}

func (c *Compiler) compileRule(rule firewall.Rule) (string, error) {
	action := "ACCEPT"
	if rule.Action == "deny" {
		action = "DROP"
	}

	proto := "tcp"
	if rule.Protocol != "" && rule.Protocol != "any" {
		proto = rule.Protocol
	}

	parts := []string{"-A", "OUTPUT"}

	if rule.Destination != "" {
		if err := c.ValidateArg(rule.Destination); err != nil {
			return "", err
		}
		parts = append(parts, "-d", rule.Destination)
	}

	parts = append(parts, "-p", proto)

	if rule.Port > 0 {
		parts = append(parts, "--dport", fmt.Sprintf("%d", rule.Port))
	}

	parts = append(parts, "-j", action)
	return strings.Join(parts, " "), nil
}
