// Package nft runs the nft command: it checks and applies rendered
// rulesets, reads counters and dynamic sets back as JSON, and maps nft's
// error messages to the objects they came from.
//
// The nft command is used rather than a netlink library: it checks a
// ruleset without applying it (nft -c), applies a whole script in one
// transaction, and its errors are readable.
package nft

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"housecarl/internal/model"
	"housecarl/internal/render"
)

// Runner runs nft.
type Runner struct {
	// Command is how nft is run: ["nft"] (the default), or a wrapper such
	// as ["docker", "run", "--rm", "-i", "--cap-add", "NET_ADMIN", "IMAGE",
	// "nft"] in tests.
	Command []string
}

// Wrapper reads HOUSECARL_NFT, a command line that runs nft (tests point
// it at a container), or returns ["nft"].
func Wrapper() []string {
	if v := strings.TrimSpace(os.Getenv("HOUSECARL_NFT")); v != "" {
		return strings.Fields(v)
	}
	return []string{"nft"}
}

func (r *Runner) cmd(ctx context.Context, args ...string) *exec.Cmd {
	c := r.Command
	if len(c) == 0 {
		c = []string{"nft"}
	}
	return exec.CommandContext(ctx, c[0], append(append([]string{}, c[1:]...), args...)...)
}

// run runs nft with args, script on stdin.
func (r *Runner) run(ctx context.Context, script string, args ...string) ([]byte, error) {
	ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()
	c := r.cmd(ctx, args...)
	if script != "" {
		c.Stdin = strings.NewReader(script)
	}
	var out, errb bytes.Buffer
	c.Stdout, c.Stderr = &out, &errb
	if err := c.Run(); err != nil {
		msg := strings.TrimSpace(errb.String())
		if msg == "" {
			msg = err.Error()
		}
		return out.Bytes(), &Error{Output: msg}
	}
	return out.Bytes(), nil
}

// Error is nft's complaint.
type Error struct {
	Output string
	// Object is what the first failing line came from, when known.
	Object string
	// Message is the first error line, without the file and position.
	Message string
}

func (e *Error) Error() string {
	if e.Object != "" && e.Message != "" {
		return e.Object + ": " + e.Message
	}
	if e.Message != "" {
		return e.Message
	}
	return e.Output
}

// errLine matches "/dev/stdin:12:3-10: Error: …".
var errLine = regexp.MustCompile(`(?m)^[^:\n]*:(\d+):[0-9-]+: Error: (.*)$`)

// explain fills in Object and Message from nft's output and the rendered
// ruleset.
func explain(err error, rs *render.Ruleset) error {
	e, ok := err.(*Error)
	if !ok {
		return err
	}
	if m := errLine.FindStringSubmatch(e.Output); m != nil {
		line, _ := strconv.Atoi(m[1])
		e.Message = strings.TrimSpace(m[2])
		e.Object = rs.Object(line)
	}
	return e
}

// Check has nft check a ruleset without applying it.
func (r *Runner) Check(ctx context.Context, rs *render.Ruleset) error {
	_, err := r.run(ctx, rs.Text, "-c", "-f", "/dev/stdin")
	return explain(err, rs)
}

// Apply applies a ruleset (one transaction: all or nothing).
func (r *Runner) Apply(ctx context.Context, rs *render.Ruleset) error {
	_, err := r.run(ctx, rs.Text, "-f", "/dev/stdin")
	return explain(err, rs)
}

// ListTable returns Housecarl's table as nft JSON (nil if it doesn't
// exist).
func (r *Runner) ListTable(ctx context.Context) ([]byte, error) {
	out, err := r.run(ctx, "", "-j", "list", "table", "inet", render.Table)
	if err != nil {
		if e, ok := err.(*Error); ok && strings.Contains(e.Output, "No such file or directory") {
			return nil, nil
		}
		return nil, err
	}
	return out, nil
}

// ListRuleset returns the whole ruleset as nft JSON (the importer's
// input on this machine).
func (r *Runner) ListRuleset(ctx context.Context) ([]byte, error) {
	return r.run(ctx, "", "-j", "list", "ruleset")
}

// TableText returns a table as an nft script that recreates it.
func (r *Runner) TableText(ctx context.Context, family, name string) (string, error) {
	out, err := r.run(ctx, "", "list", "table", family, name)
	if err != nil {
		return "", err
	}
	return "table " + family + " " + name + " {}\ndelete table " + family + " " + name + "\n" + string(out), nil
}

// ApplyScript applies an nft script (a takeover's rollback).
func (r *Runner) ApplyScript(ctx context.Context, script string) error {
	_, err := r.run(ctx, script, "-f", "/dev/stdin")
	return err
}

// DeleteTable deletes a table (hc takeover: the ruleset Housecarl
// replaces).
func (r *Runner) DeleteTable(ctx context.Context, family, name string) error {
	_, err := r.run(ctx, "", "delete", "table", family, name)
	return err
}

// AddElement adds an entry to a dynamic set's family part, with a timeout
// in seconds (0: the set's default).
func (r *Runner) AddElement(ctx context.Context, set, addr string, timeout int) error {
	a, err := model.ParseAddress(addr)
	if err != nil {
		return err
	}
	t := ""
	if timeout > 0 {
		t = fmt.Sprintf(" timeout %ds", timeout)
	}
	_, err = r.run(ctx, fmt.Sprintf("add element inet %s %s { %s%s }\n", render.Table, render.SetName(set, a.IsV6()), a, t), "-f", "/dev/stdin")
	return err
}

// DeleteElement removes an entry from a dynamic set.
func (r *Runner) DeleteElement(ctx context.Context, set, addr string) error {
	a, err := model.ParseAddress(addr)
	if err != nil {
		return err
	}
	_, err = r.run(ctx, fmt.Sprintf("delete element inet %s %s { %s }\n", render.Table, render.SetName(set, a.IsV6()), a), "-f", "/dev/stdin")
	return err
}

// ---- reading JSON

// Table is what Housecarl reads back from its table: counters by mark
// and the dynamic sets' entries.
type Table struct {
	// Counters by the comment's mark ("r12", "f3", "b-established",
	// "p-input"); several nftables rules of one object add up.
	Counters map[string]Counter
	// Elements by set name (both families together).
	Elements map[string][]render.Element
	// Quotas: the rules' quotas' used bytes (quota_r12), and the
	// per-address quotas' entries by set (quota_r12_v4).
	QuotaUsage
	// Since is when the counters started (housecarld's first apply).
	Since time.Time
	// Tables are all the machine's tables ("inet main"), when the whole
	// ruleset was read.
	Tables []string
}

// QuotaUsage is what the quotas used: kept across reloads, and saved in
// the database so a reboot doesn't reset it.
type QuotaUsage struct {
	Quotas    map[string]int64            `json:"quotas,omitempty"`
	QuotaSets map[string][]render.Element `json:"quota_sets,omitempty"`
}

// merge takes the larger usage of each quota and entry from o (what was
// saved before a reboot, when the kernel's count started again).
func (u *QuotaUsage) merge(o *QuotaUsage) {
	if o == nil {
		return
	}
	if u.Quotas == nil {
		u.Quotas = map[string]int64{}
	}
	if u.QuotaSets == nil {
		u.QuotaSets = map[string][]render.Element{}
	}
	for k, v := range o.Quotas {
		if v > u.Quotas[k] {
			u.Quotas[k] = v
		}
	}
	for set, es := range o.QuotaSets {
		for _, e := range es {
			found := false
			for i := range u.QuotaSets[set] {
				if u.QuotaSets[set][i].Addr == e.Addr {
					found = true
					if e.Used > u.QuotaSets[set][i].Used {
						u.QuotaSets[set][i].Used = e.Used
					}
				}
			}
			if !found {
				u.QuotaSets[set] = append(u.QuotaSets[set], e)
			}
		}
	}
}

// Counter is packets and bytes.
type Counter struct {
	Packets uint64 `json:"packets"`
	Bytes   uint64 `json:"bytes"`
}

type jsonDoc struct {
	Nftables []map[string]json.RawMessage `json:"nftables"`
}

type jsonRule struct {
	Comment string            `json:"comment"`
	Expr    []json.RawMessage `json:"expr"`
}

type jsonSet struct {
	Name string            `json:"name"`
	Elem []json.RawMessage `json:"elem"`
}

// jsonQuota is a quota's state: a named quota ({"bytes", "used"}, in
// bytes) or one in an element ({"val", "val_unit", "used", "used_unit"}).
type jsonQuota struct {
	Name     string `json:"name"`
	Used     int64  `json:"used"`
	UsedUnit string `json:"used_unit"`
}

func (q jsonQuota) bytes() int64 {
	if m, ok := model.ByteUnits[q.UsedUnit]; ok {
		return q.Used * m
	}
	return q.Used
}

// ParseTable reads `nft -j list table inet housecarl`.
func ParseTable(data []byte) (*Table, error) {
	t := &Table{Counters: map[string]Counter{}, Elements: map[string][]render.Element{},
		QuotaUsage: QuotaUsage{Quotas: map[string]int64{}, QuotaSets: map[string][]render.Element{}}}
	if len(data) == 0 {
		return t, nil
	}
	var doc jsonDoc
	if err := json.Unmarshal(data, &doc); err != nil {
		return nil, fmt.Errorf("reading nft's JSON: %w", err)
	}
	for _, item := range doc.Nftables {
		if raw, ok := item["rule"]; ok {
			var r jsonRule
			if json.Unmarshal(raw, &r) != nil || !strings.HasPrefix(r.Comment, "hc:") {
				continue
			}
			mark, _, _ := strings.Cut(strings.TrimPrefix(r.Comment, "hc:"), " ")
			for _, e := range r.Expr {
				var c struct {
					Counter *Counter `json:"counter"`
				}
				if json.Unmarshal(e, &c) == nil && c.Counter != nil {
					cur := t.Counters[mark]
					cur.Packets += c.Counter.Packets
					cur.Bytes += c.Counter.Bytes
					t.Counters[mark] = cur
				}
			}
		}
		if raw, ok := item["quota"]; ok {
			var q jsonQuota
			if json.Unmarshal(raw, &q) == nil && strings.HasPrefix(q.Name, "quota_r") {
				t.Quotas[q.Name] = q.bytes()
			}
		}
		if raw, ok := item["set"]; ok {
			var s jsonSet
			if json.Unmarshal(raw, &s) != nil {
				continue
			}
			if strings.HasPrefix(s.Name, "quota_r") {
				for _, e := range s.Elem {
					var w struct {
						Elem struct {
							Val   string     `json:"val"`
							Quota *jsonQuota `json:"quota"`
						} `json:"elem"`
					}
					if json.Unmarshal(e, &w) == nil && w.Elem.Quota != nil {
						t.QuotaSets[s.Name] = append(t.QuotaSets[s.Name], render.Element{Addr: w.Elem.Val, Used: w.Elem.Quota.bytes()})
					}
				}
				continue
			}
			name := strings.TrimSuffix(strings.TrimSuffix(s.Name, "_v4"), "_v6")
			if name == s.Name || strings.HasPrefix(s.Name, "meter_") || strings.HasPrefix(s.Name, "conn_") {
				continue
			}
			for _, e := range s.Elem {
				if el, ok := parseElem(e); ok {
					t.Elements[name] = append(t.Elements[name], el)
				}
			}
		}
	}
	return t, nil
}

// parseElem reads a set element: "10.0.0.5", or {"elem": {"val":
// "10.0.0.5", "timeout": 1800, "expires": 1799}}.
func parseElem(raw json.RawMessage) (render.Element, bool) {
	var s string
	if json.Unmarshal(raw, &s) == nil {
		return render.Element{Addr: s}, true
	}
	var w struct {
		Elem struct {
			Val     json.RawMessage `json:"val"`
			Expires int             `json:"expires"`
		} `json:"elem"`
	}
	if json.Unmarshal(raw, &w) != nil || json.Unmarshal(w.Elem.Val, &s) != nil {
		return render.Element{}, false
	}
	return render.Element{Addr: s, Timeout: w.Elem.Expires}, true
}

// ---- the firewall

// Firewall is the engine's Firewall: it renders configurations and has
// nft check and apply them, carrying dynamic sets' entries over.
type Firewall struct {
	Runner *Runner
	// Interfaces lists the interfaces that exist now (nil: assume all).
	Interfaces func() map[string]bool
	// BootFile is where the confirmed ruleset is saved for the boot unit
	// (housecarl-ruleset.service; "" for none).
	BootFile string

	// Replacing the table resets its counters: base keeps the totals of
	// the tables replaced since housecarld started.
	mu    sync.Mutex
	base  map[string]Counter
	since time.Time
	// restore is the quota usage saved before housecarld (or the machine)
	// stopped, put back by the next apply.
	restore *QuotaUsage
}

// RestoreQuotas has the next apply put this saved quota usage back where
// the kernel's count is lower (after a reboot, the boot ruleset starts
// every quota at zero).
func (f *Firewall) RestoreQuotas(u *QuotaUsage) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.restore = u
}

// ResetQuota starts a rule's quota again: the named quota, or the
// per-address quota sets (both families).
func (f *Firewall) ResetQuota(ctx context.Context, ruleID int64, perAddress bool) error {
	name := render.QuotaName(ruleID)
	if !perAddress {
		_, err := f.Runner.run(ctx, "", "reset", "quota", "inet", render.Table, name)
		return err
	}
	for _, fam := range []string{"_v4", "_v6"} {
		if _, err := f.Runner.run(ctx, "", "flush", "set", "inet", render.Table, name+fam); err != nil {
			if e, ok := err.(*Error); ok && strings.Contains(e.Output, "No such file or directory") {
				continue // the rule has no such family
			}
			return err
		}
	}
	return nil
}

func (f *Firewall) options() render.Options {
	var o render.Options
	if f.Interfaces != nil {
		o.Interfaces = f.Interfaces()
	}
	return o
}

// Render builds the ruleset for st as Apply would.
func (f *Firewall) Render(st *model.State) *render.Ruleset { return render.Render(st, f.options()) }

// Check implements engine.Firewall.
func (f *Firewall) Check(ctx context.Context, st *model.State) error {
	return f.Runner.Check(ctx, render.Render(st, f.options()))
}

// Apply implements engine.Firewall.
func (f *Firewall) Apply(ctx context.Context, st *model.State) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	opt := f.options()
	var old *Table
	if data, err := f.Runner.ListTable(ctx); err == nil {
		if t, err := ParseTable(data); err == nil {
			opt.Elements = t.Elements
			old = t
		}
	}
	usage := QuotaUsage{}
	if old != nil {
		usage = old.QuotaUsage
	}
	usage.merge(f.restore)
	opt.Quotas, opt.QuotaSets = usage.Quotas, usage.QuotaSets
	if err := f.Runner.Apply(ctx, render.Render(st, opt)); err != nil {
		return err
	}
	f.restore = nil
	if f.base == nil {
		f.base, f.since = map[string]Counter{}, time.Now()
	}
	if old != nil {
		for k, c := range old.Counters {
			b := f.base[k]
			b.Packets += c.Packets
			b.Bytes += c.Bytes
			f.base[k] = b
		}
	}
	return nil
}

// SaveBoot writes the ruleset of st for the boot unit, atomically.
func (f *Firewall) SaveBoot(st *model.State) error {
	if f.BootFile == "" {
		return nil
	}
	rs := render.Render(st, f.options())
	tmp := f.BootFile + ".tmp"
	if err := os.MkdirAll(filepath.Dir(f.BootFile), 0o750); err != nil {
		return err
	}
	if err := os.WriteFile(tmp, []byte(rs.Text), 0o640); err != nil {
		return err
	}
	return os.Rename(tmp, f.BootFile)
}

// Table reads Housecarl's table back, with counters totalled since
// housecarld started (across reloads).
func (f *Firewall) Table(ctx context.Context) (*Table, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	data, err := f.Runner.ListTable(ctx)
	if err != nil {
		return nil, err
	}
	if data == nil {
		return nil, ErrNoTable
	}
	t, err := ParseTable(data)
	if err != nil {
		return nil, err
	}
	for k, b := range f.base {
		c := t.Counters[k]
		c.Packets += b.Packets
		c.Bytes += b.Bytes
		t.Counters[k] = c
	}
	t.Since = f.since
	return t, nil
}

// ErrNoTable: Housecarl's table isn't in the ruleset.
var ErrNoTable = errors.New("there is no table inet housecarl: Housecarl's ruleset isn't in force")

// AddElement adds an entry to a dynamic set.
func (f *Firewall) AddElement(ctx context.Context, set, addr string, timeout int) error {
	return f.Runner.AddElement(ctx, set, addr, timeout)
}

// DeleteElement removes an entry from a dynamic set.
func (f *Firewall) DeleteElement(ctx context.Context, set, addr string) error {
	return f.Runner.DeleteElement(ctx, set, addr)
}

// ListRuleset returns the machine's whole ruleset as JSON.
func (f *Firewall) ListRuleset(ctx context.Context) ([]byte, error) { return f.Runner.ListRuleset(ctx) }

// TableText returns another table as a script that recreates it.
func (f *Firewall) TableText(ctx context.Context, family, name string) (string, error) {
	return f.Runner.TableText(ctx, family, name)
}

// ApplyScript applies an nft script.
func (f *Firewall) ApplyScript(ctx context.Context, script string) error {
	return f.Runner.ApplyScript(ctx, script)
}

// DeleteTable deletes another table (takeover).
func (f *Firewall) DeleteTable(ctx context.Context, family, name string) error {
	if name == render.Table && family == "inet" {
		return fmt.Errorf("that is Housecarl's own table")
	}
	return f.Runner.DeleteTable(ctx, family, name)
}