package firewall

import (
	"fmt"
	"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{}
	table := ""
	for _, line := range strings.Split(live, "\n") {
		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+" ")
				if strings.HasSuffix(line, "-j "+original) {
					line = strings.TrimSuffix(line, "-j "+original) + "-j " + owned
				}
			}
			out.WriteString(line + "\n")
		}
		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()
}