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: inserting PREROUTING DNAT rule: src=%s proto=%s sport=%s -> dst=%s dport=%s comment=%q",
		sourceIP, proto, sourcePort, targetIP, targetPort, comment)
	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 {
	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(sourceCIDR, proto, sourcePort, comment string) error {
	logger.Info("IPTABLES: inserting POSTROUTING MASQUERADE rule: src=%s proto=%s sport=%s comment=%q",
		sourceCIDR, proto, sourcePort, comment)
	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 {
	logger.Info("IPTABLES: inserting POSTROUTING MASQUERADE rule for target: dst=%s proto=%s dport=%s comment=%q",
		targetCIDR, proto, targetPort, comment)
	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 {
	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)
	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 {
	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
}

// 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)
	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 {
	logger.Info("IPTABLES: inserting POSTROUTING MASQUERADE rule in container PID %d: src=%s proto=%s sport=%s comment=%q",
		pid, sourceCIDR, proto, sourcePort, comment)
	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...)
}