package firewall

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

func ApplyDualStack(filters []FilterRule, nats []NATRule, mangles []MangleRule) error {
	for _, r := range filters {
		if err := validateRule("filter", r); err != nil {
			return err
		}
	}
	for _, r := range nats {
		if err := validateRule("nat", r); err != nil {
			return err
		}
	}
	for _, r := range mangles {
		if err := validateRule("mangle", r); err != nil {
			return err
		}
	}
	var f4, f6 []FilterRule
	var n4, n6 []NATRule
	var m4, m6 []MangleRule
	validate := func(family int, srcGeo, dstGeo string) error {
		if family != 0 && family != 4 && family != 6 {
			return fmt.Errorf("IP family must be 4 or 6")
		}
		if family == 6 && (srcGeo != "" || dstGeo != "") {
			return fmt.Errorf("GeoIP datasets currently support IPv4 only; remove GeoIP matching from the IPv6 rule")
		}
		return nil
	}
	for _, r := range filters {
		if err := validate(r.IPVersion, r.SrcGeoIP, r.DstGeoIP); err != nil {
			return err
		}
		if r.IPVersion == 6 {
			if r.Protocol == "icmp" {
				r.Protocol = "ipv6-icmp"
			}
			f6 = append(f6, r)
		} else {
			f4 = append(f4, r)
		}
	}
	for _, r := range nats {
		if err := validate(r.IPVersion, r.SrcGeoIP, r.DstGeoIP); err != nil {
			return err
		}
		if r.IPVersion == 6 {
			n6 = append(n6, r)
		} else {
			n4 = append(n4, r)
		}
	}
	for _, r := range mangles {
		if err := validate(r.IPVersion, r.SrcGeoIP, r.DstGeoIP); err != nil {
			return err
		}
		if r.IPVersion == 6 {
			m6 = append(m6, r)
		} else {
			m4 = append(m4, r)
		}
	}
	type family struct {
		binary  string
		payload string
	}
	families := []family{{"iptables", BuildIPTablesRestore(f4, n4, m4)}}
	six := BuildIPTablesRestore(f6, n6, m6)
	sets := regexp.MustCompile(`(--(?:match|add)-set) ([^ ]+)`)
	six = sets.ReplaceAllString(six, "${1} ${2}_v6")
	live6, err := exec.Command("ip6tables-save").Output()
	need6 := len(f6)+len(n6)+len(m6) > 0 || strings.Contains(string(live6), "DT_")
	if need6 {
		if err != nil {
			return fmt.Errorf("IPv6 firewall unavailable: %w", err)
		}
		families = append(families, family{"ip6tables", six})
	}
	// Validate all families before applying either one.
	for i := range families {
		live, err := exec.Command(families[i].binary+"-save", "-c").Output()
		if err != nil {
			return err
		}
		families[i].payload = ownedPayload(families[i].payload, string(live))
		cmd := exec.Command(families[i].binary+"-restore", "--noflush", "-c", "--test")
		cmd.Stdin = strings.NewReader(families[i].payload)
		if out, err := cmd.CombinedOutput(); err != nil {
			return fmt.Errorf("%s validation: %s: %w", families[i].binary, out, err)
		}
	}
	for _, family := range families {
		cmd := exec.Command(family.binary+"-restore", "--noflush", "-c")
		cmd.Stdin = strings.NewReader(family.payload)
		if out, err := cmd.CombinedOutput(); err != nil {
			return fmt.Errorf("%s apply: %s: %w", family.binary, out, err)
		}
	}
	return nil
}