package blocking

import (
	"fmt"
	"math"
	"net/netip"
	"sort"
	"strings"
)

// Client steps (docs/blocking.md): named checks policy zones in order and
// the first match decides, so a "gate" zone that lets some clients through
// (rpz-passthru, matched on the client's address) keeps them from every
// zone after it. Jarl puts a gate in front of the lists of each step,
// letting through the clients whose step is below it.

// GateSuffix follows a gate's name: off.gate.jarl.invalid,
// strict.gate.jarl.invalid.
const GateSuffix = ".gate.jarl.invalid"

// v4Mapped holds the IPv4-mapped IPv6 addresses.
var v4Mapped = netip.MustParsePrefix("::ffff:0:0/96")

// GateZone returns the policy zone name of the gate in front of a step's
// lists ("off" for the gate in front of the normal lists).
func GateZone(step string) string { return step + GateSuffix }

// Assignment gives the addresses in Prefix a step rank (-1 off, 0 normal,
// 1, 2, ... the steps above it).
type Assignment struct {
	Prefix netip.Prefix
	Rank   int
}

// PassBelow returns the fewest prefixes that cover exactly the addresses whose
// rank is below k. An address has the rank of the most specific
// assignment containing it, and rank 0 (normal) when none does. When two
// assignments have the same prefix, the higher rank counts (validation
// refuses that anyway).
func PassBelow(assigns []Assignment, k int) []netip.Prefix {
	var out []netip.Prefix
	for _, root := range []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0"), netip.MustParsePrefix("::/0")} {
		var mine []Assignment
		if !root.Addr().Is4() {
			// named matches IPv4 clients as IPv4-mapped IPv6 addresses:
			// keep IPv6 prefixes clear of them, so only the IPv4 ones
			// decide for IPv4 clients.
			mine = append(mine, Assignment{v4Mapped, math.MaxInt})
		}
		rank := 0
		for _, a := range assigns {
			p := a.Prefix.Masked()
			if p.Addr().Is4() != root.Addr().Is4() {
				continue
			}
			if p.Bits() == 0 {
				rank = max(rank, a.Rank)
				continue
			}
			mine = append(mine, Assignment{p, a.Rank})
		}
		for _, p := range cover(root, rank, mine, k) {
			if p.Bits() == 0 {
				// A trigger needs at least one bit: use both halves.
				lo, hi := halves(p)
				out = append(out, lo, hi)
				continue
			}
			out = append(out, p)
		}
	}
	return out
}

// cover returns the prefixes inside p whose addresses have a rank below k.
// rank is p's own; inside are the assignments strictly inside p.
func cover(p netip.Prefix, rank int, inside []Assignment, k int) []netip.Prefix {
	if len(inside) == 0 {
		if rank < k {
			return []netip.Prefix{p}
		}
		return nil
	}
	var out []netip.Prefix
	lo, hi := halves(p)
	full := 0
	for _, h := range []netip.Prefix{lo, hi} {
		var sub []Assignment
		for _, a := range inside {
			if a.Prefix.Bits() > h.Bits() && h.Contains(a.Prefix.Addr()) {
				sub = append(sub, a)
			}
		}
		c := cover(h, ownRank(inside, h, rank), sub, k)
		if len(c) == 1 && c[0] == h {
			full++
		}
		out = append(out, c...)
	}
	if full == 2 {
		return []netip.Prefix{p}
	}
	return out
}

// ownRank is the rank of h: its own assignments' (the highest), or the
// inherited one without any.
func ownRank(inside []Assignment, h netip.Prefix, inherited int) int {
	r, found := 0, false
	for _, a := range inside {
		if a.Prefix == h {
			if !found || a.Rank > r {
				r = a.Rank
			}
			found = true
		}
	}
	if !found {
		return inherited
	}
	return r
}

// halves splits a prefix into its two halves.
func halves(p netip.Prefix) (netip.Prefix, netip.Prefix) {
	bits := p.Bits() + 1
	lo := netip.PrefixFrom(p.Addr(), bits).Masked()
	b := lo.Addr().AsSlice()
	i := (bits - 1) / 8
	b[i] |= 0x80 >> ((bits - 1) % 8)
	a, _ := netip.AddrFromSlice(b)
	return lo, netip.PrefixFrom(a, bits)
}

// ClientIPOwner returns the policy zone owner name that matches clients
// in p (an rpz-client-ip trigger): 32.37.0.10.10.rpz-client-ip for
// 10.10.0.37/32, 64.zz.3d7f.8013.2a02.rpz-client-ip for 2a02:8013:3d7f::/64.
func ClientIPOwner(p netip.Prefix) string {
	p = p.Masked()
	a := p.Addr()
	var labels []string
	if a.Is4() {
		b := a.As4()
		labels = []string{fmt.Sprint(b[3]), fmt.Sprint(b[2]), fmt.Sprint(b[1]), fmt.Sprint(b[0])}
	} else {
		// RFC 5952 text, with "zz" for the "::" and the words reversed.
		s := a.String()
		var words []string
		if left, right, ok := strings.Cut(s, "::"); ok {
			if left != "" {
				words = append(words, strings.Split(left, ":")...)
			}
			words = append(words, "zz")
			if right != "" {
				words = append(words, strings.Split(right, ":")...)
			}
		} else {
			words = strings.Split(s, ":")
		}
		for i := len(words) - 1; i >= 0; i-- {
			labels = append(labels, words[i])
		}
	}
	return fmt.Sprint(p.Bits()) + "." + strings.Join(labels, ".") + ".rpz-client-ip"
}

// GateZoneFile renders a gate: every client in prefixes passes.
func GateZoneFile(title string, prefixes []netip.Prefix) string {
	owners := make([]string, 0, len(prefixes))
	for _, p := range prefixes {
		owners = append(owners, ClientIPOwner(p))
	}
	sort.Strings(owners)
	var sb strings.Builder
	sb.WriteString(header(title, 1))
	for _, o := range owners {
		sb.WriteString(o + "\tCNAME\trpz-passthru.\n")
	}
	return sb.String()
}