package firewall

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

// ownedPayload confines DebTik rules to dedicated chains. The host's chains,
// policies and rules are left intact; only legacy DebTik ID rules are migrated.
func ownedPayload(payload, live string) string {
	liveTables := map[string][]string{}
	counts := map[string][2]uint64{}
	oldChains := map[string][]string{}
	counted := map[string]bool{}
	counterPattern := regexp.MustCompile(`^\[(\d+):(\d+)\] (.*)$`)
	commentPattern := regexp.MustCompile(`--comment (?:"id:([^"]+)"|id:([^ ]+))`)
	table := ""
	for _, line := range strings.Split(live, "\n") {
		var pkts, bytes uint64
		if m := counterPattern.FindStringSubmatch(line); m != nil {
			pkts, _ = strconv.ParseUint(m[1], 10, 64)
			bytes, _ = strconv.ParseUint(m[2], 10, 64)
			line = m[3]
		}
		if m := commentPattern.FindStringSubmatch(line); m != nil {
			id := m[1] + m[2]
			key := table + ":" + id
			c := counts[key]
			counts[key] = [2]uint64{c[0] + pkts, c[1] + bytes}
		}
		if strings.HasPrefix(line, ":DT_") {
			oldChains[table] = append(oldChains[table], strings.Fields(line)[0][1:])
		}
		if strings.HasPrefix(line, "*") {
			table = line[1:]
		} else if strings.HasPrefix(line, "-A ") {
			liveTables[table] = append(liveTables[table], line)
		}
	}
	var out strings.Builder
	for _, block := range strings.Split(payload, "COMMIT") {
		lines := strings.Split(strings.TrimSpace(block), "\n")
		if len(lines) == 0 || !strings.HasPrefix(lines[0], "*") {
			continue
		}
		table = lines[0][1:]
		names := map[string]string{}
		var chains []string
		for _, line := range lines[1:] {
			if strings.HasPrefix(line, ":") {
				name := strings.Fields(line)[0][1:]
				names[name] = "DT_" + table + "_" + name
				chains = append(chains, name)
			}
		}
		fmt.Fprintf(&out, "*%s\n", table)
		for _, chain := range chains {
			fmt.Fprintf(&out, ":%s - [0:0]\n-F %s\n", names[chain], names[chain])
		}
		for _, line := range liveTables[table] {
			if !strings.HasPrefix(strings.Fields(line)[1], "DT_") && strings.Contains(line, "--comment") && (strings.Contains(line, "id:filter-") || strings.Contains(line, "id:nat-") || strings.Contains(line, "id:mangle-")) {
				out.WriteString("-D " + strings.TrimPrefix(line, "-A ") + "\n")
			}
		}
		for _, line := range lines[1:] {
			fields := strings.Fields(line)
			if len(fields) < 2 || fields[0] != "-A" {
				continue
			}
			line = strings.Replace(line, "-A "+fields[1]+" ", "-A "+names[fields[1]]+" ", 1)
			for original, owned := range names {
				line = strings.ReplaceAll(line, "-j "+original+" ", "-j "+owned+" ")
				line = strings.ReplaceAll(line, "-g "+original+" ", "-g "+owned+" ")
				if strings.HasSuffix(line, "-g "+original) {
					line = strings.TrimSuffix(line, "-g "+original) + "-g " + owned
				}
				if strings.HasSuffix(line, "-j "+original) {
					line = strings.TrimSuffix(line, "-j "+original) + "-j " + owned
				}
			}
			if m := commentPattern.FindStringSubmatch(line); m != nil {
				key := table + ":" + m[1] + m[2]
				if !counted[key] {
					c := counts[key]
					if c[0] != 0 || c[1] != 0 {
						line = fmt.Sprintf("[%d:%d] %s", c[0], c[1], line)
					}
					counted[key] = true
				}
			}
			out.WriteString(line + "\n")
		}
		wanted := map[string]bool{}
		for _, name := range names {
			wanted[name] = true
		}
		var obsolete []string
		for _, old := range oldChains[table] {
			if wanted[old] {
				continue
			}
			foreign := false
			for _, rule := range liveTables[table] {
				source := strings.Fields(rule)[1]
				if !strings.HasPrefix(source, "DT_") && (strings.HasSuffix(rule, "-j "+old) || strings.HasSuffix(rule, "-g "+old)) {
					foreign = true
				}
			}
			if !foreign {
				obsolete = append(obsolete, old)
			}
		}
		for _, old := range obsolete {
			fmt.Fprintf(&out, "-F %s\n", old)
		}
		for _, old := range obsolete {
			fmt.Fprintf(&out, "-X %s\n", old)
		}
		for _, chain := range chains {
			builtin := chain == "INPUT" || chain == "OUTPUT" || chain == "FORWARD" || chain == "PREROUTING" || chain == "POSTROUTING"
			if !builtin {
				continue
			}
			exists := false
			for _, line := range liveTables[table] {
				if strings.HasPrefix(line, "-A "+chain+" ") && strings.Contains(line, "debtik:hook") && strings.HasSuffix(line, "-j "+names[chain]) {
					exists = true
					break
				}
			}
			if !exists {
				fmt.Fprintf(&out, "-I %s 1 -m comment --comment debtik:hook -j %s\n", chain, names[chain])
			}
		}
		out.WriteString("COMMIT\n")
	}
	return out.String()
}