package iptables

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

// IPTablesAPI defines the interface for iptables operations, enabling mock implementations for testing
type IPTablesAPI interface {
	Binary() string
	EnsureIPForward() error
	EnsureEstablishedRelated(chain string) error
	DeleteLine(chain string, lineNum string) error
	InsertPreroutingRule(sourceIP, proto, sourcePort, targetIP, targetPort, comment string) error
	InsertPreroutingRuleOnInterface(iface, proto, sourcePort, targetIP, targetPort, comment string) error
	InsertPostroutingMasquerade(sourceCIDR, proto, sourcePort, comment string) error
	InsertPostroutingMasqueradeForTarget(targetCIDR, proto, targetPort, comment string) error
	InsertForwardAccept(chain, sourceIP, targetIP, proto, sourcePort, targetPort, comment string) error
	DeleteForwardAccept(chain, comment string) error
	InsertPreroutingRuleInContainer(pid int, sourceIP, proto, sourcePort, targetIP, targetPort, comment string) error
	InsertPostroutingMasqueradeInContainer(pid int, sourceCIDR, proto, sourcePort, comment string) error
}

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

// Ensure Manager implements IPTablesAPI
var _ IPTablesAPI = (*Manager)(nil)

// 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 := "/usr/sbin/iptables-legacy"
	if !strings.Contains(m.binary, "legacy") {
		iptPath = "/usr/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.
// Logs a warning if it fails (e.g. read-only filesystem in a container),
// since this should be configured at the host level.
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
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 {
		return nil
	}

	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
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...)
	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 {
	iptPath := "/usr/sbin/iptables-legacy"
	if !strings.Contains(m.binary, "legacy") {
		iptPath = "/usr/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])
			}
		}
	}

	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 {
	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)
	}

	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", "-t", "nat", "-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 {
	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 {
	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)
	}

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

	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 {
	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...)
}