package iptables

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

	"firewall_containers/network-go/logger"
)

// 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() {
	logger.Info("IPTABLES: detecting iptables binary")
	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"
		logger.Info("IPTABLES: detected iptables-legacy (DOCKER-USER chain present)")
		return
	}
	m.binary = "/usr/sbin/iptables"
	logger.Info("IPTABLES: using default iptables binary")
}

// 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 {
	cmdStr := m.binary + " " + strings.Join(args, " ")
	logger.Info("IPTABLES: executing: %s", cmdStr)
	if m.debug {
		logger.Debug("IPTABLES DEBUG: %s %s", m.binary, strings.Join(args, " "))
	}
	cmd := exec.Command(m.binary, args...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		logger.Error("IPTABLES: command failed: %s\noutput: %s", cmdStr, strings.TrimSpace(string(output)))
		return fmt.Errorf("iptables %s failed: %w\noutput: %s", strings.Join(args, " "), err, string(output))
	}
	logger.Debug("IPTABLES: command succeeded: %s", cmdStr)
	if len(output) > 0 {
		logger.Debug("IPTABLES: output: %s", strings.TrimSpace(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 {
	// Use the same binary path that was auto-detected on the host
	// Alpine installs iptables at /usr/sbin/ not /sbin/
	iptPath := m.binary

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

	cmdStr := "nsenter " + strings.Join(fullArgs, " ")
	logger.Info("IPTABLES: executing nsenter for PID %d: %s", pid, cmdStr)
	if m.debug {
		logger.Debug("IPTABLES DEBUG: nsenter %s", strings.Join(fullArgs, " "))
	}

	cmd := exec.Command("nsenter", fullArgs...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		logger.Error("IPTABLES: nsenter command failed for PID %d\noutput: %s", pid, strings.TrimSpace(string(output)))
		return fmt.Errorf("nsenter iptables failed: %w\noutput: %s", err, string(output))
	}
	logger.Info("IPTABLES: nsenter command succeeded for PID %d", pid)
	if len(output) > 0 {
		logger.Debug("IPTABLES: nsenter output: %s", strings.TrimSpace(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 {
	logger.Info("IPTABLES: enabling IP forwarding (echo 1 > /proc/sys/net/ipv4/ip_forward)")
	cmd := exec.Command("sh", "-c", "echo 1 > /proc/sys/net/ipv4/ip_forward")
	output, err := cmd.CombinedOutput()
	if err != nil {
		logger.Warn("IPTABLES: failed to enable ip_forward (expected if read-only fs): %v", err)
		return fmt.Errorf("failed to enable ip_forward: %w\noutput: %s", err, string(output))
	}
	logger.Info("IPTABLES: IP forwarding enabled")
	return nil
}

// EnsureEstablishedRelated inserts an ESTABLISHED,RELATED accept rule at the top of a chain
func (m *Manager) EnsureEstablishedRelated(chain string) error {
	logger.Debug("IPTABLES: checking for ESTABLISHED,RELATED rule in %s", chain)
	checkArgs := []string{"-w", "-n", "-L", chain}
	cmd := exec.Command(m.binary, checkArgs...)
	output, err := cmd.Output()
	if err != nil {
		logger.Debug("IPTABLES: could not list chain %s: %v", chain, err)
		return nil
	}

	if !strings.Contains(string(output), "ESTABLISHED") || !strings.Contains(string(output), "RELATED") {
		logger.Info("IPTABLES: inserting ESTABLISHED,RELATED ACCEPT rule in %s", chain)
		args := []string{"-w", "-I", chain, "-m", "state", "--state", "established,related", "-j", "ACCEPT"}
		return m.run(args...)
	}
	logger.Debug("IPTABLES: ESTABLISHED,RELATED rule already exists in %s", chain)
	return nil
}

// DeleteLine deletes a specific line number from a chain
func (m *Manager) DeleteLine(chain string, lineNum string) error {
	logger.Info("IPTABLES: deleting line %s from chain %s", lineNum, chain)
	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 {
	logger.Info("IPTABLES: deleting line %s from chain %s in container PID %d", lineNum, chain, pid)
	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 {
		logger.Debug("IPTABLES: getLineNumbers failed for %s: %v", chain, err)
		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])
			}
		}
	}
	logger.Debug("IPTABLES: getLineNumbers chain=%s patterns=%v found=%v", chain, grepPatterns, matchingLines)
	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...)
	if len(lines) > 0 {
		logger.Info("IPTABLES: deleting %d matching lines from %s: %v", len(lines), chain, lines)
	}
	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 := m.binary

	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 {
		logger.Debug("IPTABLES: deleteMatchingLinesInContainer list failed for PID %d chain %s: %v", pid, chain, err)
		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])
			}
		}
	}

	if len(matchingLines) > 0 {
		logger.Info("IPTABLES: deleting %d matching lines from container PID %d chain %s: %v",
			len(matchingLines), pid, chain, matchingLines)
	}

	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 {
	logger.Info("IPTABLES: checking PREROUTING DNAT rule: src=%s proto=%s sport=%s -> dst=%s dport=%s comment=%q",
		sourceIP, proto, sourcePort, targetIP, targetPort, comment)

	// Idempotent: check if rule already exists
	existing := m.getLineNumbers("PREROUTING", "nat", comment, "DNAT", targetIP, targetPort)
	if len(existing) > 0 {
		logger.Debug("IPTABLES: PREROUTING DNAT rule already exists (lines=%v), skipping", existing)
		return nil
	}

	logger.Info("IPTABLES: inserting PREROUTING DNAT rule: src=%s proto=%s sport=%s -> dst=%s dport=%s comment=%q",
		sourceIP, proto, sourcePort, targetIP, targetPort, comment)
	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 {
	logger.Info("IPTABLES: checking PREROUTING DNAT rule on interface %s: proto=%s dport=%s -> %s:%s comment=%q",
		iface, proto, sourcePort, targetIP, targetPort, comment)

	// Check if rule already exists (idempotent: don't re-apply)
	// Must include port in check to distinguish port 80 from port 443 rules with same comment
	existing := m.getLineNumbers("PREROUTING", "nat", comment, "DNAT", targetIP, targetPort)
	if len(existing) > 0 {
		logger.Debug("IPTABLES: PREROUTING DNAT rule already exists on %s (lines=%v), skipping", iface, existing)
		return nil
	}

	logger.Info("IPTABLES: inserting PREROUTING DNAT rule on interface %s: proto=%s dport=%s -> %s:%s comment=%q",
		iface, proto, sourcePort, targetIP, targetPort, comment)
	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(destCIDR, proto, destPort, comment string) error {
	logger.Info("IPTABLES: checking POSTROUTING MASQUERADE rule: dst=%s proto=%s dport=%s comment=%q",
		destCIDR, proto, destPort, comment)

	// Check if rule already exists (idempotent: don't re-apply)
	// Must include port in check to distinguish port 80 from port 443 rules with same comment
	existing := m.getLineNumbers("POSTROUTING", "nat", comment, "MASQUERADE", destCIDR, destPort)
	if len(existing) > 0 {
		logger.Debug("IPTABLES: POSTROUTING MASQUERADE rule already exists (lines=%v), skipping", existing)
		return nil
	}

	logger.Info("IPTABLES: inserting POSTROUTING MASQUERADE rule: dst=%s proto=%s dport=%s comment=%q",
		destCIDR, proto, destPort, comment)
	args := []string{
		"-w", "-t", "nat", "-I", "POSTROUTING",
		"-d", destCIDR,
		"-p", proto,
		"--dport", destPort,
		"-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 {
	logger.Info("IPTABLES: checking POSTROUTING MASQUERADE rule for target: dst=%s proto=%s dport=%s comment=%q",
		targetCIDR, proto, targetPort, comment)

	// Idempotent: check if rule already exists
	// Must include port in check to distinguish port 80 from port 443 rules with same comment
	existing := m.getLineNumbers("POSTROUTING", "nat", comment, "MASQUERADE", targetCIDR, targetPort)
	if len(existing) > 0 {
		logger.Debug("IPTABLES: POSTROUTING MASQUERADE for target already exists (lines=%v), skipping", existing)
		return nil
	}

	logger.Info("IPTABLES: inserting POSTROUTING MASQUERADE rule for target: dst=%s proto=%s dport=%s comment=%q",
		targetCIDR, proto, targetPort, comment)
	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 {
	logger.Info("IPTABLES: checking FORWARD ACCEPT rule: chain=%s src=%s dst=%s proto=%s sport=%s dport=%s comment=%q",
		chain, sourceIP, targetIP, proto, sourcePort, targetPort, comment)

	// Idempotent: check if rule already exists
	var grepPatterns []string
	grepPatterns = append(grepPatterns, comment, 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)
	}
	existing := m.getLineNumbers(chain, "", grepPatterns...)
	if len(existing) > 0 {
		logger.Debug("IPTABLES: FORWARD ACCEPT rule already exists in %s (lines=%v), skipping", chain, existing)
		return nil
	}

	logger.Info("IPTABLES: inserting FORWARD ACCEPT rule: chain=%s src=%s dst=%s proto=%s sport=%s dport=%s comment=%q",
		chain, sourceIP, targetIP, proto, sourcePort, targetPort, comment)
	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 {
	logger.Info("IPTABLES: deleting FORWARD ACCEPT rules in %s with comment=%q", chain, comment)
	lines := m.getLineNumbers(chain, "", comment)
	if len(lines) > 0 {
		logger.Info("IPTABLES: found %d FORWARD ACCEPT rules to delete: %v", len(lines), lines)
	}
	for i := len(lines) - 1; i >= 0; i-- {
		if err := m.DeleteLine(chain, lines[i]); err != nil {
			return err
		}
	}
	return nil
}

// checkContainerChainExists lists rules inside a container namespace and returns
// the output, or an error if nsenter/iptables fails inside the container.
func (m *Manager) checkContainerChainExists(pid int, table, chain string) (string, error) {
	iptPath := m.binary
	nsenterArgs := []string{"-t", fmt.Sprintf("%d", pid), "-n", "--", iptPath, "-w", "-n", "-t", table, "-L", chain}
	cmd := exec.Command("nsenter", nsenterArgs...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("nsenter iptables list failed for PID %d chain %s/%s: %w (output: %s)", pid, table, chain, err, strings.TrimSpace(string(output)))
	}
	return string(output), 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 {
	logger.Info("IPTABLES: inserting PREROUTING DNAT rule in container PID %d: src=%s proto=%s dport=%s -> %s:%s comment=%q",
		pid, sourceIP, proto, sourcePort, targetIP, targetPort, comment)

	// First, try to list the chain inside the container to check state
	output, err := m.checkContainerChainExists(pid, "nat", "PREROUTING")
	if err != nil {
		// Cannot inspect container iptables — return error so caller falls back to host
		logger.Warn("IPTABLES: cannot check container PREROUTING (PID %d): %v", pid, err)
		return fmt.Errorf("cannot check container iptables: %w", err)
	}

	// Idempotent check: scan existing rules for matching patterns
	ruleExists := false
	for _, line := range strings.Split(output, "\n") {
		if strings.Contains(line, "DNAT") &&
			strings.Contains(line, sourcePort) &&
			strings.Contains(line, targetIP) &&
			strings.Contains(line, comment) {
			ruleExists = true
			break
		}
	}
	if ruleExists {
		logger.Info("IPTABLES: PREROUTING DNAT rule already exists in container PID %d (dport=%s -> %s), skipping", pid, sourcePort, targetIP)
		return nil
	}

	// Rule doesn't exist — clean up stale/duplicate rules then insert
	patterns := []string{"DNAT", sourcePort, targetIP, targetPort, comment}
	if delErr := m.deleteMatchingLinesInContainer(pid, "nat", "PREROUTING", patterns...); delErr != nil {
		logger.Debug("IPTABLES: stale PREROUTING cleanup in container PID %d: %v", pid, delErr)
	}

	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, destCIDR, proto, destPort, comment string) error {
	logger.Info("IPTABLES: inserting POSTROUTING MASQUERADE rule in container PID %d: dst=%s proto=%s dport=%s comment=%q",
		pid, destCIDR, proto, destPort, comment)

	// First, try to list the chain inside the container to check state
	output, err := m.checkContainerChainExists(pid, "nat", "POSTROUTING")
	if err != nil {
		logger.Warn("IPTABLES: cannot check container POSTROUTING (PID %d): %v", pid, err)
		return fmt.Errorf("cannot check container iptables: %w", err)
	}

	// Idempotent check: scan existing rules for matching patterns
	// Must include port to distinguish port 80 from port 443 rules with same comment
	ruleExists := false
	for _, line := range strings.Split(output, "\n") {
		if strings.Contains(line, "MASQUERADE") &&
			strings.Contains(line, comment) &&
			strings.Contains(line, destCIDR) &&
			strings.Contains(line, destPort) {
			ruleExists = true
			break
		}
	}
	if ruleExists {
		logger.Info("IPTABLES: POSTROUTING MASQUERADE rule already exists in container PID %d (dst=%s dport=%s), skipping", pid, destCIDR, destPort)
		return nil
	}

	// Rule doesn't exist — clean up stale/duplicate rules then insert
	patterns := []string{"MASQUERADE", comment, destCIDR, destPort}
	if delErr := m.deleteMatchingLinesInContainer(pid, "nat", "POSTROUTING", patterns...); delErr != nil {
		logger.Debug("IPTABLES: stale POSTROUTING cleanup in container PID %d: %v", pid, delErr)
	}

	args := []string{
		"-I", "POSTROUTING",
		"-d", destCIDR,
		"-p", proto,
		"--dport", destPort,
		"-m", "comment", "--comment", comment,
		"-j", "MASQUERADE",
	}
	return m.runInContainer(pid, "nat", args...)
}