package lb

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net"
	"os/exec"
	"strings"

	"github.com/coreos/go-iptables/iptables"
	"github.com/vishvananda/netlink"
	"golang.org/x/sys/unix"
)

const (
	// ForwardChain (nat, from PREROUTING) DNATs held addresses' ports to endpoints.
	ForwardChain = "REWEAVE-LB"
	// PostChain (nat, from POSTROUTING, ahead of weave's WEAVE) keeps the client's address on what ForwardChain DNATed.
	PostChain = "REWEAVE-LB-POST"
	// MarkChain (mangle, from PREROUTING) marks holders' connections and steers their replies.
	MarkChain = "REWEAVE-LB-MARK"
	// servicePrefix names the per-port chains ForwardChain jumps to.
	servicePrefix = "REWEAVE-LBS-"

	markShift = 24
	markMask  = uint32(0xff) << markShift
	// RulePriority is where the reply-steering rules sit: before the main table (32766).
	RulePriority = 9000
	// TableBase + mark is the routing table that sends replies to that mark's holder.
	TableBase = 52000

	bridge = "weave"
)

func markValue(m uint32) uint32 { return m << markShift }

// portChain is the per-(service, protocol, port) chain name: stable, short enough for iptables (28 chars).
func portChain(f Forward, p Port) string {
	h := sha256.Sum256([]byte(fmt.Sprintf("%s/%s/%s/%d", f.Namespace, f.Name, p.Protocol, p.Port)))
	return servicePrefix + strings.ToUpper(hex.EncodeToString(h[:])[:12])
}

// Rules is the full content of the chains this package owns, in iptables-restore form.
type Rules struct {
	Nat, Mangle, Filter string
}

// Render turns the desired state (load-balancer forwards and egress) into iptables-restore input for the nat,
// mangle and filter tables. podNet is the weave allocation range (pods and bridges).
// Each payload declares every chain it owns, so a restore with --noflush replaces exactly those.
func Render(d Desired, e EgressDesired, self string, podNet *net.IPNet, stale []string) Rules {
	var nat bytes.Buffer
	nat.WriteString("*nat\n")
	for _, c := range []string{ForwardChain, PostChain, ClassChain} {
		fmt.Fprintf(&nat, ":%s - [0:0]\n", c)
	}
	var body bytes.Buffer
	seen := map[string]bool{}
	for _, f := range d.Forwards {
		a := f.Address.String()
		if seen[a] {
			continue
		}
		seen[a] = true
		// weave masquerades everything entering the pod network from outside unless it is for a local pod
		// (its expose rule); a connection we DNATed must keep the client's address, so it leaves the nat
		// table here. One from a pod (hairpin) is masqueraded to this node, so its reply comes back here.
		fmt.Fprintf(&body, "-A %s -m conntrack --ctstate DNAT --ctorigdst %s/32 -s %s -j MASQUERADE\n", PostChain, a, podNet)
		fmt.Fprintf(&body, "-A %s -m conntrack --ctstate DNAT --ctorigdst %s/32 -j ACCEPT\n", PostChain, a)
	}
	for _, f := range d.Forwards {
		for _, p := range f.Ports {
			c := portChain(f, p)
			fmt.Fprintf(&nat, ":%s - [0:0]\n", c)
			proto := strings.ToLower(string(p.Protocol))
			fmt.Fprintf(&body, "-A %s -d %s/32 -p %s -m %s --dport %d -m comment --comment %q -j %s\n",
				ForwardChain, f.Address, proto, proto, p.Port, f.Namespace+"/"+f.Name, c)
			n := len(p.Endpoints)
			for i, e := range p.Endpoints {
				// DNAT to an ip:port needs the protocol in the same rule
				if i < n-1 {
					// the i-th endpoint takes 1/(n-i) of what the earlier ones left: uniform overall
					fmt.Fprintf(&body, "-A %s -p %s -m statistic --mode random --probability %.10f -j DNAT --to-destination %s\n",
						c, proto, 1/float64(n-i), e)
				} else {
					fmt.Fprintf(&body, "-A %s -p %s -j DNAT --to-destination %s\n", c, proto, e)
				}
			}
		}
	}
	for _, c := range stale {
		// chains of ports no longer forwarded: emptied here, deleted after the restore
		fmt.Fprintf(&nat, ":%s - [0:0]\n", c)
	}

	var mangle bytes.Buffer
	mangle.WriteString("*mangle\n")
	for _, c := range []string{MarkChain, MSSChain} {
		fmt.Fprintf(&mangle, ":%s - [0:0]\n", c)
	}
	if len(d.Holders) > 0 {
		// replies of a holder's connection carry its mark: the policy rule routes them back to it
		fmt.Fprintf(&mangle, "-A %s -m conntrack --ctdir REPLY -m connmark ! --mark 0x0/0x%x -j CONNMARK --restore-mark --nfmask 0x%x --ctmask 0x%x\n",
			MarkChain, markMask, markMask, markMask)
		for _, h := range d.Holders {
			// a new connection to a pod here, arriving from the holder's bridge, from outside the pod network:
			// the holder DNATed it and must see its replies
			fmt.Fprintf(&mangle, "-A %s -i %s -m mac --mac-source %s ! -s %s -d %s -m conntrack --ctstate NEW -m comment --comment %q -j CONNMARK --set-xmark 0x%x/0x%x\n",
				MarkChain, bridge, h.Node.WeaveMAC, podNet, podNet, "holder "+h.Node.Name, markValue(h.Mark), markMask)
		}
	}

	var filter bytes.Buffer
	filter.WriteString("*filter\n")
	for _, c := range []string{FwdChain, InChain} {
		fmt.Fprintf(&filter, ":%s - [0:0]\n", c)
	}
	renderEgress(e, self, podNet, &body, &mangle, &filter)

	nat.Write(body.Bytes())
	nat.WriteString("COMMIT\n")
	mangle.WriteString("COMMIT\n")
	filter.WriteString("COMMIT\n")
	return Rules{Nat: nat.String(), Mangle: mangle.String(), Filter: filter.String()}
}

// Datapath applies Desired to this node.
type Datapath struct {
	ipt *iptables.IPTables
	// restore runs iptables-restore --noflush with the given input (swappable in tests)
	restore func(input string) error
}

// NewDatapath uses the host's iptables (whichever backend the image's launch script selected).
func NewDatapath(ipt *iptables.IPTables) *Datapath {
	return &Datapath{ipt: ipt, restore: func(input string) error {
		cmd := exec.Command("iptables-restore", "--noflush", "-w", "5")
		cmd.Stdin = strings.NewReader(input)
		if out, err := cmd.CombinedOutput(); err != nil {
			return fmt.Errorf("iptables-restore: %v: %s", err, out)
		}
		return nil
	}}
}

// Apply makes the node's iptables chains, tunnel, policy rules and routes exactly what d and e say.
func (dp *Datapath) Apply(d Desired, e EgressDesired, self string, local Local) error {
	podNet := local.PodNet
	if err := dp.hooks(); err != nil {
		return err
	}
	// the tunnels must exist before the routes into them (iptables rules may name a device that does not exist yet)
	var peers []Node
	if e.Tunnelled(self) {
		peers = e.Peers(self)
	}
	if err := ensureTunnels(peers, local); err != nil {
		return err
	}
	existing, err := dp.ipt.ListChains("nat")
	if err != nil {
		return err
	}
	want := map[string]bool{}
	for _, f := range d.Forwards {
		for _, p := range f.Ports {
			want[portChain(f, p)] = true
		}
	}
	var stale []string
	for _, c := range existing {
		if strings.HasPrefix(c, servicePrefix) && !want[c] {
			stale = append(stale, c)
		}
	}
	r := Render(d, e, self, podNet, stale)
	// filter first: the guards are in place before the marks and SNATs that they guard
	for _, payload := range []string{r.Filter, r.Nat, r.Mangle} {
		if err := dp.restore(payload); err != nil {
			return err
		}
	}
	for _, c := range stale {
		if err := dp.ipt.DeleteChain("nat", c); err != nil {
			return fmt.Errorf("deleting stale chain %s: %v", c, err)
		}
	}
	if err := dp.routes(d.Holders); err != nil {
		return err
	}
	return egressNet(e, self)
}

// hooks makes sure the built-in chains jump to ours: first (kube-proxy's and weave's rules must not see these
// connections first), except ClassChain, which must come after kube-proxy's KUBE-SERVICES (whose DNAT ends the
// nat walk, so a connection to a Service is never classed as going outside).
func (dp *Datapath) hooks() error {
	for _, h := range []struct {
		table, from, chain string
		last               bool
	}{
		{"nat", "PREROUTING", ForwardChain, false},
		{"nat", "POSTROUTING", PostChain, false},
		{"nat", "PREROUTING", ClassChain, true},
		{"mangle", "PREROUTING", MarkChain, false},
		{"mangle", "FORWARD", MSSChain, false},
		{"filter", "FORWARD", FwdChain, false},
		{"filter", "INPUT", InChain, false},
	} {
		ok, err := dp.ipt.ChainExists(h.table, h.chain)
		if err != nil {
			return err
		}
		if !ok {
			if err := dp.ipt.NewChain(h.table, h.chain); err != nil {
				return err
			}
		}
		rule := []string{"-m", "comment", "--comment", "reweave lb", "-j", h.chain}
		rules, err := dp.ipt.List(h.table, h.from)
		if err != nil {
			return err
		}
		if h.last {
			ours, kube := -1, -1
			for i, r := range rules {
				if strings.HasSuffix(r, "-j "+h.chain) {
					ours = i
				}
				if strings.HasSuffix(r, "-j KUBE-SERVICES") {
					kube = i
				}
			}
			if ours > kube {
				continue
			}
			if err := dp.ipt.DeleteIfExists(h.table, h.from, rule...); err != nil {
				return err
			}
			if err := dp.ipt.Append(h.table, h.from, rule...); err != nil {
				return err
			}
			continue
		}
		// rules[0] is the policy line; our jump must be rules[1]
		if len(rules) > 1 && strings.HasSuffix(rules[1], "-j "+h.chain) {
			continue
		}
		if err := dp.ipt.DeleteIfExists(h.table, h.from, rule...); err != nil {
			return err
		}
		if err := dp.ipt.Insert(h.table, h.from, 1, rule...); err != nil {
			return err
		}
	}
	return nil
}

// routes keeps one policy rule and one table per holder: marked replies go to the holder's bridge.
func (dp *Datapath) routes(holders []Holder) error {
	link, err := netlink.LinkByName(bridge)
	if err != nil {
		if len(holders) == 0 {
			return nil
		}
		return fmt.Errorf("weave bridge: %v", err)
	}
	want := map[int]Holder{}
	for _, h := range holders {
		want[TableBase+int(h.Mark)] = h
	}
	rules, err := netlink.RuleList(unix.AF_INET)
	if err != nil {
		return err
	}
	have := map[int]bool{}
	for _, r := range rules {
		if r.Priority != RulePriority || r.Table < TableBase || r.Table > TableBase+MaxMarks {
			continue
		}
		h, ok := want[r.Table]
		if ok && uint32(r.Mark) == markValue(h.Mark) && uint32(r.Mask) == markMask {
			have[r.Table] = true
			continue
		}
		r := r
		if err := netlink.RuleDel(&r); err != nil {
			return fmt.Errorf("removing stale rule to table %d: %v", r.Table, err)
		}
	}
	existing, err := routesByTable(TableBase+1, TableBase+MaxMarks)
	if err != nil {
		return err
	}
	for table := TableBase + 1; table <= TableBase+MaxMarks; table++ {
		h, ok := want[table]
		if !ok {
			for _, rt := range existing[table] {
				if err := delRoute(rt); err != nil {
					return fmt.Errorf("removing stale route in table %d: %v", table, err)
				}
			}
			continue
		}
		if err := netlink.RouteReplace(&netlink.Route{
			LinkIndex: link.Attrs().Index,
			Dst:       &net.IPNet{IP: net.IPv4zero, Mask: net.CIDRMask(0, 32)},
			Gw:        h.Node.WeaveIP,
			Table:     table,
			Flags:     int(netlink.FLAG_ONLINK),
		}); err != nil {
			return fmt.Errorf("route to holder %s: %v", h.Node.Name, err)
		}
		if !have[table] {
			rule := netlink.NewRule()
			rule.Priority = RulePriority
			rule.Table = table
			rule.Mark = int(markValue(h.Mark))
			rule.Mask = int(markMask)
			if err := netlink.RuleAdd(rule); err != nil {
				return fmt.Errorf("rule for holder %s: %v", h.Node.Name, err)
			}
		}
	}
	return nil
}

// delRoute deletes a route as RouteList returned it: a default route lists without Dst, which RouteDel refuses.
func delRoute(rt netlink.Route) error {
	if rt.Dst == nil {
		rt.Dst = &net.IPNet{IP: net.IPv4zero, Mask: net.CIDRMask(0, 32)}
	}
	return netlink.RouteDel(&rt)
}