package iptables

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

// Manager manages iptables rules via the iptables/iptables-legacy CLI
type Manager struct {
	binary string // /usr/sbin/iptables or /usr/sbin/iptables-legacy
	debug  bool
}

// NewManager creates a new iptables manager, auto-detecting the binary
func NewManager(debug bool) *Manager {
	m := &Manager{debug: debug}
	m.detectBinary()
	return m
}

// detectBinary checks if iptables-legacy is available (matching shell script logic)
func (m *Manager) detectBinary() {
	cmd := exec.Command("/usr/sbin/iptables-legacy", "-L")
	output, err := cmd.CombinedOutput()
	if err == nil && strings.Contains(string(output), "DOCKER-USER") {
		m.binary = "/usr/sbin/iptables-legacy"
		return
	}
	m.binary = "/usr/sbin/iptables"
}

// Binary returns the detected iptables binary path
func (m *Manager) Binary() string {
	return m.binary
}

// run executes an iptables command on the host
func (m *Manager) run(args ...string) error {
	if m.debug {
		fmt.Printf("[IPTABLES DEBUG] %s %s\n", m.binary, strings.Join(args, " "))
	}
	cmd := exec.Command(m.binary, args...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("iptables %s failed: %w\noutput: %s", strings.Join(args, " "), err, string(output))
	}
	return nil
}

// runInContainer executes an iptables command inside a container's network namespace via nsenter
func (m *Manager) runInContainer(pid int, table string, args ...string) error {
	iptPath := "/sbin/iptables-legacy"
	if !strings.Contains(m.binary, "legacy") {
		iptPath = "/sbin/iptables"
	}

	fullArgs := []string{"-t", fmt.Sprintf("%d", pid), "-n", "--", iptPath}
	if table != "" {
		fullArgs = append(fullArgs, "-t", table)
	}
	fullArgs = append(fullArgs, args...)

	if m.debug {
		fmt.Printf("[IPTABLES DEBUG] nsenter %s\n", strings.Join(fullArgs, " "))
	}
	cmd := exec.Command("nsenter", fullArgs...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("nsenter iptables failed: %w\noutput: %s", err, string(output))
	}
	return nil
}

// EnsureIPForward enables IP forwarding on the host
func (m *Manager) EnsureIPForward() error {
	cmd := exec.Command("sh", "-c", "echo 1 > /proc/sys/net/ipv4/ip_forward")
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to enable ip_forward: %w\noutput: %s", err, string(output))
	}
	return nil
}

// EnsureEstablishedRelated inserts an ESTABLISHED,RELATED accept rule at the top of a chain
// if it doesn't already exist
func (m *Manager) EnsureEstablishedRelated(chain string) error {
	checkArgs := []string{"-w", "-n", "-L", chain}
	cmd := exec.Command(m.binary, checkArgs...)
	output, err := cmd.Output()
	if err != nil {
		// Chain may not exist, create it
		return nil
	}

	// Only insert if ESTABLISHED,RELATED rule is not present
	if !strings.Contains(string(output), "ESTABLISHED") || !strings.Contains(string(output), "RELATED") {
		args := []string{"-w", "-I", chain, "-m", "state", "--state", "established,related", "-j", "ACCEPT"}
		return m.run(args...)
	}
	return nil
}

// DeleteLine deletes a specific line number from a chain
func (m *Manager) DeleteLine(chain string, lineNum string) error {
	args := []string{"-w", "-D", chain, lineNum}
	return m.run(args...)
}

// DeleteLineInContainer deletes a specific line number from a chain inside a container namespace
func (m *Manager) DeleteLineInContainer(pid int, table, chain, lineNum string) error {
	args := []string{"-D", chain, lineNum}
	return m.runInContainer(pid, table, args...)
}

// getLineNumbers returns line numbers matching certain criteria in a chain/table
// This implements the grep logic from the shell script: iptables -w --line-number -n -L $CHAIN | grep ...
func (m *Manager) getLineNumbers(chain, table string, grepPatterns ...string) []string {
	args := []string{"-w", "--line-number", "-n", "-L", chain}
	if table != "" {
		args = []string{"-w", "-t", table, "--line-number", "-n", "-L", chain}
	}

	cmd := exec.Command(m.binary, args...)
	output, err := cmd.Output()
	if err != nil {
		return nil
	}

	lines := strings.Split(string(output), "\n")
	var matchingLines []string
	for _, line := range lines {
		matchesAll := true
		for _, pattern := range grepPatterns {
			if !strings.Contains(line, pattern) {
				matchesAll = false
				break
			}
		}
		if matchesAll {
			fields := strings.Fields(line)
			if len(fields) > 0 {
				matchingLines = append(matchingLines, fields[0])
			}
		}
	}
	return matchingLines
}

// deleteMatchingLines deletes all lines in a chain matching the given patterns
func (m *Manager) deleteMatchingLines(chain, table string, grepPatterns ...string) error {
	lines := m.getLineNumbers(chain, table, grepPatterns...)
	// Reverse order (highest line first) so deletions don't shift line numbers
	for i := len(lines) - 1; i >= 0; i-- {
		if err := m.DeleteLine(chain, lines[i]); err != nil {
			return err
		}
	}
	return nil
}

// deleteMatchingLinesInContainer deletes matching lines inside a container namespace
func (m *Manager) deleteMatchingLinesInContainer(pid int, table, chain string, grepPatterns ...string) error {
	// For container namespaces, we use a different approach: list via nsenter + grep
	iptPath := "/sbin/iptables-legacy"
	if !strings.Contains(m.binary, "legacy") {
		iptPath = "/sbin/iptables"
	}

	nsenterArgs := []string{"-t", fmt.Sprintf("%d", pid), "-n", "--", iptPath, "-w", "--line-number", "-n", "-t", table, "-L", chain}
	cmd := exec.Command("nsenter", nsenterArgs...)
	output, err := cmd.Output()
	if err != nil {
		return nil
	}

	lines := strings.Split(string(output), "\n")
	var matchingLines []string
	for _, line := range lines {
		matchesAll := true
		for _, pattern := range grepPatterns {
			if !strings.Contains(line, pattern) {
				matchesAll = false
				break
			}
		}
		if matchesAll {
			fields := strings.Fields(line)
			if len(fields) > 0 {
				matchingLines = append(matchingLines, fields[0])
			}
		}
	}

	// Delete in reverse order
	for i := len(matchingLines) - 1; i >= 0; i-- {
		if err := m.DeleteLineInContainer(pid, table, chain, matchingLines[i]); err != nil {
			return err
		}
	}
	return nil
}

// InsertPreroutingRule inserts a DNAT PREROUTING rule on the host
func (m *Manager) InsertPreroutingRule(sourceIP, proto, sourcePort, targetIP, targetPort, comment string) error {
	// First, delete existing matching rules
	patterns := []string{"DNAT", sourcePort, targetIP, targetPort, comment}
	if err := m.deleteMatchingLines("PREROUTING", "nat", patterns...); err != nil {
		return fmt.Errorf("failed to delete old PREROUTING rules: %w", err)
	}

	// Insert the new rule
	args := []string{
		"-w", "-t", "nat", "-I", "PREROUTING",
		"-d", sourceIP,
		"-p", proto,
		"--dport", sourcePort,
		"-m", "comment", "--comment", comment,
		"-j", "DNAT", "--to", targetIP + ":" + targetPort,
	}
	return m.run(args...)
}

// InsertPreroutingRuleOnInterface inserts a DNAT PREROUTING rule on a specific interface
func (m *Manager) InsertPreroutingRuleOnInterface(iface, proto, sourcePort, targetIP, targetPort, comment string) error {
	args := []string{
		"-w", "-I", "PREROUTING",
		"-i", iface,
		"-p", proto,
		"--dport", sourcePort,
		"-m", "comment", "--comment", comment,
		"-j", "DNAT", "--to", targetIP + ":" + targetPort,
	}
	return m.run(args...)
}

// InsertPostroutingMasquerade inserts a MASQUERADE POSTROUTING rule on the host
func (m *Manager) InsertPostroutingMasquerade(sourceCIDR, proto, sourcePort, comment string) error {
	// Delete existing matching rules first
	patterns := []string{"MASQUERADE", comment, sourceCIDR, sourcePort}
	if err := m.deleteMatchingLines("POSTROUTING", "nat", patterns...); err != nil {
		return fmt.Errorf("failed to delete old POSTROUTING rules: %w", err)
	}

	args := []string{
		"-w", "-t", "nat", "-I", "POSTROUTING",
		"-s", sourceCIDR,
		"-p", proto,
		"--sport", sourcePort,
		"-m", "comment", "--comment", comment,
		"-j", "MASQUERADE",
	}
	return m.run(args...)
}

// InsertPostroutingMasqueradeForTarget inserts a MASQUERADE POSTROUTING rule for a target
func (m *Manager) InsertPostroutingMasqueradeForTarget(targetCIDR, proto, targetPort, comment string) error {
	patterns := []string{"MASQUERADE", comment, targetCIDR, targetPort}
	if err := m.deleteMatchingLines("POSTROUTING", "nat", patterns...); err != nil {
		return fmt.Errorf("failed to delete old POSTROUTING rules: %w", err)
	}

	args := []string{
		"-w", "-t", "nat", "-I", "POSTROUTING",
		"-d", targetCIDR,
		"-p", proto,
		"--dport", targetPort,
		"-m", "comment", "--comment", comment,
		"-j", "MASQUERADE",
	}
	return m.run(args...)
}

// InsertForwardAccept inserts a FORWARD ACCEPT rule on the host
func (m *Manager) InsertForwardAccept(chain, sourceIP, targetIP, proto, sourcePort, targetPort, comment string) error {
	// Build grep patterns to match existing rules
	var grepPatterns []string
	grepPatterns = append(grepPatterns, proto)
	if sourceIP != "" {
		grepPatterns = append(grepPatterns, sourceIP)
	}
	if targetIP != "" {
		grepPatterns = append(grepPatterns, targetIP)
	}
	if sourcePort != "" {
		grepPatterns = append(grepPatterns, sourcePort)
	}
	if targetPort != "" {
		grepPatterns = append(grepPatterns, targetPort)
	}

	// Delete old matching rules
	if err := m.deleteMatchingLines(chain, "", grepPatterns...); err != nil {
		return fmt.Errorf("failed to delete old FORWARD rules: %w", err)
	}

	// Build iptables args
	args := []string{"-w", "-I", chain, "-p", proto}
	if sourceIP != "" {
		args = append(args, "-s", sourceIP)
	}
	if targetIP != "" {
		args = append(args, "-d", targetIP)
	}
	if sourcePort != "" {
		args = append(args, "--sport", sourcePort)
	}
	if targetPort != "" {
		args = append(args, "--dport", targetPort)
	}
	args = append(args, "-m", "comment", "--comment", comment, "-j", "ACCEPT")

	return m.run(args...)
}

// DeleteForwardAccept deletes a FORWARD ACCEPT rule by comment
func (m *Manager) DeleteForwardAccept(chain, comment string) error {
	lines := m.getLineNumbers(chain, "", comment)
	for i := len(lines) - 1; i >= 0; i-- {
		if err := m.DeleteLine(chain, lines[i]); err != nil {
			return err
		}
	}
	return nil
}

// InsertPreroutingRuleInContainer inserts a DNAT PREROUTING rule inside a container namespace
func (m *Manager) InsertPreroutingRuleInContainer(pid int, sourceIP, proto, sourcePort, targetIP, targetPort, comment string) error {
	// Delete existing first
	patterns := []string{"DNAT", sourcePort, targetIP, targetPort, comment}
	if err := m.deleteMatchingLinesInContainer(pid, "nat", "PREROUTING", patterns...); err != nil {
		return fmt.Errorf("failed to delete old container PREROUTING rules: %w", err)
	}

	args := []string{
		"-I", "PREROUTING",
		"-d", sourceIP,
		"-p", proto,
		"--dport", sourcePort,
		"-m", "comment", "--comment", comment,
		"-j", "DNAT", "--to", targetIP + ":" + targetPort,
	}
	return m.runInContainer(pid, "nat", args...)
}

// InsertPostroutingMasqueradeInContainer inserts a MASQUERADE POSTROUTING rule inside a container namespace
func (m *Manager) InsertPostroutingMasqueradeInContainer(pid int, sourceCIDR, proto, sourcePort, comment string) error {
	patterns := []string{"MASQUERADE", comment, sourceCIDR, sourcePort}
	if err := m.deleteMatchingLinesInContainer(pid, "nat", "POSTROUTING", patterns...); err != nil {
		return fmt.Errorf("failed to delete old container POSTROUTING rules: %w", err)
	}

	args := []string{
		"-I", "POSTROUTING",
		"-s", sourceCIDR,
		"-p", proto,
		"--sport", sourcePort,
		"-m", "comment", "--comment", comment,
		"-j", "MASQUERADE",
	}
	return m.runInContainer(pid, "nat", args...)
}