// SPDX-License-Identifier: MIT

// Package helper implements the privileged core invoked by the unprivileged
// prx CLI through pkexec.
//
// Every plan is re-validated here before any system change is made, and the
// apply operation is transactional: on any failure the changes made so far are
// rolled back. Concurrent apply/remove operations are serialised with an
// advisory file lock.
package helper

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"prx/internal/config"
	"prx/internal/engine"
	"prx/internal/execx"
	"prx/internal/fsx"
	"prx/internal/health"
	"prx/internal/iprule"
	"prx/internal/netdetect"
	"prx/internal/nft"
	"prx/internal/proto"
	"prx/internal/state"
)

// healthTimeout bounds the post-apply connectivity probe.
const healthTimeout = 3 * time.Second

// Execute runs a single privileged request.
func Execute(ctx context.Context, r execx.Runner, plan proto.Plan) proto.Result {
	switch plan.Action {
	case proto.ActionApply:
		return apply(ctx, r, plan)
	case proto.ActionRemove:
		return remove(ctx, r, plan)
	case proto.ActionInspect:
		return inspect(ctx, r, plan)
	case proto.ActionEngine:
		return engineAction(ctx, r, plan)
	case proto.ActionInitEngine:
		return initEngine(plan)
	default:
		return proto.Result{Action: plan.Action, Error: fmt.Sprintf("unknown action %q", plan.Action)}
	}
}

// rollback records undo actions in the order they must be reversed.
type rollback struct {
	fns []func(context.Context) error
}

func (rb *rollback) add(fn func(context.Context) error) { rb.fns = append(rb.fns, fn) }

// run executes the undo actions in reverse order and returns every failure so
// the caller can report that rollback was incomplete.
func (rb *rollback) run(ctx context.Context) []error {
	var errs []error
	for i := len(rb.fns) - 1; i >= 0; i-- {
		if err := rb.fns[i](ctx); err != nil {
			errs = append(errs, err)
		}
	}
	return errs
}

func apply(ctx context.Context, r execx.Runner, plan proto.Plan) (res proto.Result) {
	res = proto.Result{Action: proto.ActionApply}
	if err := validatePlan(plan); err != nil {
		res.Error = err.Error()
		return res
	}
	if err := preflight(r, plan.ManageEngine && !plan.SkipEngine && plan.ServiceName != ""); err != nil {
		res.Error = err.Error()
		return res
	}

	// A dry run makes no changes, so it does not need the state lock.
	if !plan.DryRun {
		unlock, err := lockState(plan.StatePath)
		if err != nil {
			res.Error = err.Error()
			return res
		}
		defer func() {
			if uerr := unlock(); uerr != nil {
				res.Messages = append(res.Messages, "warning: "+uerr.Error())
			}
		}()
	}

	wan, err := netdetect.Detect(ctx, r, plan.WANInterface, plan.WANGateway4, plan.WANGateway6)
	if err != nil {
		res.Error = err.Error()
		return res
	}
	if wan.Gateway4 == "" && wan.Gateway6 == "" {
		res.Error = "no default gateway detected; cannot install policy routing"
		return res
	}

	script, err := nft.Script(nft.Params{
		WANInterface: wan.Interface,
		Mark:         plan.Mark,
		Local4:       wan.Local4,
		Local6:       wan.Local6,
	})
	if err != nil {
		res.Error = err.Error()
		return res
	}

	if plan.DryRun {
		if err := nft.Check(ctx, r, script); err != nil {
			res.Error = "dry-run validation failed: " + err.Error()
			return res
		}
		res.OK = true
		res.Messages = append(res.Messages, "would protect replies: "+describeTarget(wan, plan.Mark))
		return res
	}

	rb := &rollback{}
	changed := false

	// fail aborts the transaction: it undoes everything applied so far and
	// reports any rollback step that itself failed, so nothing is ever lost
	// silently.
	fail := func(msg string) proto.Result {
		res.Error = msg
		var failed int
		for _, rbErr := range rb.run(ctx) {
			failed++
			res.Messages = append(res.Messages, "rollback warning: "+rbErr.Error())
		}
		if failed > 0 {
			res.Error = fmt.Sprintf("%s (rollback incomplete: %d step(s) failed — inspect with 'prx doctor')", msg, failed)
		}
		return res
	}

	// 1) Protect inbound replies before the engine takes over routing.
	if err := nft.Apply(ctx, r, script); err != nil {
		return fail("apply nftables rules: " + err.Error())
	}
	rb.add(func(c context.Context) error { return nft.Remove(c, r) })
	changed = true

	// 2) Policy routing for marked packets. The rollback for a table is
	// registered before the table is touched so that a failure part-way
	// through still removes anything that was added.
	for _, target := range []struct {
		family   int
		gateway  string
		table    int
		priority int
	}{
		{family: 4, gateway: wan.Gateway4, table: plan.TableV4, priority: plan.PriorityV4},
		{family: 6, gateway: wan.Gateway6, table: plan.TableV6, priority: plan.PriorityV6},
	} {
		if target.gateway == "" {
			continue
		}
		family, table := target.family, target.table
		rb.add(func(c context.Context) error { return removeTable(c, r, family, table) })
		spec := iprule.Spec{
			Family:   family,
			Mark:     plan.Mark,
			Table:    table,
			Priority: target.priority,
			Gateway:  target.gateway,
			Dev:      wan.Interface,
		}
		if err := ensureRouting(ctx, r, spec, &changed); err != nil {
			return fail(err.Error())
		}
	}

	// 3) Destination-based bypass (direct) routing.
	if err := applyBypass(ctx, r, plan, wan, rb, &changed); err != nil {
		return fail(err.Error())
	}

	// 4) Ensure the engine is running.
	engineWasRunning := false
	if plan.ManageEngine && !plan.SkipEngine && plan.ServiceName != "" {
		active, err := engine.IsActive(ctx, r, plan.ServiceName)
		if err != nil {
			return fail(err.Error())
		}
		engineWasRunning = active
		if !active {
			if err := engine.Start(ctx, r, plan.ServiceName); err != nil {
				return fail(err.Error())
			}
			changed = true
			service := plan.ServiceName
			rb.add(func(c context.Context) error { return engine.Stop(c, r, service) })
		}
	}

	// 5) Health checks.
	checks, fatalOK := runChecks(ctx, r, plan, wan)
	res.Checks = checks
	if !fatalOK {
		return fail("post-apply health checks failed; changes rolled back")
	}

	// 6) Record state.
	st := &state.State{
		Enabled:          true,
		WANInterface:     wan.Interface,
		Gateway4:         wan.Gateway4,
		Gateway6:         wan.Gateway6,
		Mark:             plan.Mark,
		TableV4:          plan.TableV4,
		TableV6:          plan.TableV6,
		PriorityV4:       plan.PriorityV4,
		PriorityV6:       plan.PriorityV6,
		BypassTableV4:    plan.BypassTableV4,
		BypassTableV6:    plan.BypassTableV6,
		BypassPrioV4:     plan.BypassPrioV4,
		BypassPrioV6:     plan.BypassPrioV6,
		BypassCIDRs:      plan.BypassCIDRs,
		EngineManaged:    plan.ManageEngine && !plan.SkipEngine,
		EngineWasRunning: engineWasRunning,
		EnabledAt:        time.Now().UTC(),
	}
	if err := state.Save(plan.StatePath, st); err != nil {
		return fail(err.Error())
	}

	res.OK = true
	res.Changed = changed
	res.State = &proto.StateInfo{
		Enabled:      true,
		WANInterface: wan.Interface,
		Gateway4:     wan.Gateway4,
		Gateway6:     wan.Gateway6,
		Mark:         plan.Mark,
	}
	res.Messages = append(res.Messages, "proxy routing enabled: "+describeTarget(wan, plan.Mark))
	return res
}

// describeTarget renders a human readable summary of the applied configuration.
func describeTarget(wan netdetect.WAN, mark uint32) string {
	parts := []string{"interface " + wan.Interface}
	if wan.Gateway4 != "" {
		parts = append(parts, "gw4 "+wan.Gateway4)
	}
	if wan.Gateway6 != "" {
		parts = append(parts, "gw6 "+wan.Gateway6)
	}
	parts = append(parts, fmt.Sprintf("mark 0x%x", mark))
	return strings.Join(parts, ", ")
}

func ensureRouting(ctx context.Context, r execx.Runner, spec iprule.Spec, changed *bool) error {
	ch, err := iprule.EnsureRule(ctx, r, spec)
	if err != nil {
		return err
	}
	*changed = *changed || ch

	ch, err = iprule.EnsureRoute(ctx, r, spec.Family, spec.Table, spec.Gateway, spec.Dev)
	if err != nil {
		return err
	}
	*changed = *changed || ch
	return nil
}

func applyBypass(ctx context.Context, r execx.Runner, plan proto.Plan, wan netdetect.WAN, rb *rollback, changed *bool) error {
	registeredV4, registeredV6 := false, false
	for _, cidr := range plan.BypassCIDRs {
		ip, _, err := net.ParseCIDR(cidr)
		if err != nil {
			return fmt.Errorf("invalid bypass CIDR %q: %w", cidr, err)
		}

		family, table, prio, gw := 4, plan.BypassTableV4, plan.BypassPrioV4, wan.Gateway4
		if ip.To4() == nil {
			family, table, prio, gw = 6, plan.BypassTableV6, plan.BypassPrioV6, wan.Gateway6
		}
		if gw == "" {
			// No gateway for this family on this host; skip silently.
			continue
		}

		// Register the table cleanup before touching it.
		if family == 4 && !registeredV4 {
			t := table
			rb.add(func(c context.Context) error { return removeTable(c, r, 4, t) })
			registeredV4 = true
		}
		if family == 6 && !registeredV6 {
			t := table
			rb.add(func(c context.Context) error { return removeTable(c, r, 6, t) })
			registeredV6 = true
		}

		if ch, err := iprule.EnsureCIDRRoute(ctx, r, family, table, cidr, gw, wan.Interface); err != nil {
			return err
		} else {
			*changed = *changed || ch
		}
		spec := iprule.BypassSpec{Family: family, CIDR: cidr, Table: table, Priority: prio, Gateway: gw, Dev: wan.Interface}
		if ch, err := iprule.EnsureBypassRule(ctx, r, spec); err != nil {
			return err
		} else {
			*changed = *changed || ch
		}
	}
	return nil
}

func removeTable(ctx context.Context, r execx.Runner, family, table int) error {
	if _, err := iprule.RemoveRulesByTable(ctx, r, family, table); err != nil {
		return err
	}
	return iprule.FlushTable(ctx, r, family, table)
}

func remove(ctx context.Context, r execx.Runner, plan proto.Plan) (res proto.Result) {
	res = proto.Result{Action: proto.ActionRemove}
	if err := validatePlan(plan); err != nil {
		res.Error = err.Error()
		return res
	}
	if err := preflight(r, plan.StopEngineOnOff && plan.ManageEngine && plan.ServiceName != ""); err != nil {
		res.Error = err.Error()
		return res
	}

	unlock, err := lockState(plan.StatePath)
	if err != nil {
		res.Error = err.Error()
		return res
	}
	defer func() {
		if uerr := unlock(); uerr != nil {
			res.Messages = append(res.Messages, "warning: "+uerr.Error())
		}
	}()

	var msgs []string
	for _, family := range []int{4, 6} {
		table := plan.TableV4
		bp := plan.BypassTableV4
		if family == 6 {
			table = plan.TableV6
			bp = plan.BypassTableV6
		}
		for _, spec := range []struct {
			label string
			table int
		}{
			{label: fmt.Sprintf("ipv%d policy", family), table: table},
			{label: fmt.Sprintf("ipv%d bypass", family), table: bp},
		} {
			n, err := iprule.RemoveRulesByTable(ctx, r, family, spec.table)
			if err != nil {
				res.Error = err.Error()
				return res
			}
			if n > 0 {
				msgs = append(msgs, fmt.Sprintf("removed %d %s rule(s)", n, spec.label))
			}
			if err := iprule.FlushTable(ctx, r, family, spec.table); err != nil {
				res.Error = err.Error()
				return res
			}
		}
	}

	if err := nft.Remove(ctx, r); err != nil {
		res.Error = "remove nftables table: " + err.Error()
		return res
	}

	if plan.StopEngineOnOff && plan.ManageEngine && plan.ServiceName != "" {
		if err := engine.Stop(ctx, r, plan.ServiceName); err != nil {
			res.Messages = append(res.Messages, "warning: "+err.Error())
		} else {
			msgs = append(msgs, "stopped engine "+plan.ServiceName)
		}
	}

	if err := state.Remove(plan.StatePath); err != nil {
		res.Messages = append(res.Messages, "warning: "+err.Error())
	}

	res.OK = true
	res.Changed = true
	res.Messages = append(msgs, "proxy routing disabled")
	res.State = &proto.StateInfo{Enabled: false}
	return res
}

func inspect(ctx context.Context, r execx.Runner, plan proto.Plan) proto.Result {
	res := proto.Result{Action: proto.ActionInspect, OK: true}

	exists, err := nft.Exists(ctx, r)
	if err != nil {
		res.Checks = append(res.Checks, proto.Check{Name: "nft-table", OK: false, Fatal: true, Detail: err.Error()})
	} else {
		res.Checks = append(res.Checks, proto.Check{Name: "nft-table", OK: exists, Fatal: true, Detail: nft.Table()})
	}

	wan, err := netdetect.Detect(ctx, r, plan.WANInterface, plan.WANGateway4, plan.WANGateway6)
	if err != nil {
		res.Checks = append(res.Checks, proto.Check{Name: "wan", OK: false, Detail: err.Error(), Fatal: true})
		return res
	}
	res.Checks = append(res.Checks, proto.Check{Name: "wan", OK: true, Detail: describeTarget(wan, plan.Mark)})

	for _, target := range []struct {
		family   int
		gateway  string
		table    int
		priority int
	}{
		{family: 4, gateway: wan.Gateway4, table: plan.TableV4, priority: plan.PriorityV4},
		{family: 6, gateway: wan.Gateway6, table: plan.TableV6, priority: plan.PriorityV6},
	} {
		if target.gateway == "" {
			// No gateway for this family on this host: nothing to check.
			continue
		}
		spec := iprule.Spec{Family: target.family, Mark: plan.Mark, Table: target.table, Priority: target.priority}
		ruleOK, err := iprule.RuleExists(ctx, r, spec)
		res.Checks = append(res.Checks, proto.Check{
			Name:   fmt.Sprintf("ipv%d-rule", target.family),
			OK:     ruleOK && err == nil,
			Fatal:  true,
			Detail: errDetail(err, fmt.Sprintf("fwmark 0x%x -> table %d prio %d", plan.Mark, target.table, target.priority)),
		})
		routeOK, err := iprule.TableHasRoutes(ctx, r, target.family, target.table)
		res.Checks = append(res.Checks, proto.Check{
			Name:   fmt.Sprintf("ipv%d-route", target.family),
			OK:     routeOK && err == nil,
			Fatal:  true,
			Detail: errDetail(err, "table "+strconv.Itoa(target.table)),
		})
	}

	if plan.ManageEngine && plan.ServiceName != "" {
		active, err := engine.IsActive(ctx, r, plan.ServiceName)
		res.Checks = append(res.Checks, proto.Check{Name: "engine", OK: active && err == nil, Warn: true, Detail: errDetail(err, plan.ServiceName)})
	}

	tuns, tunErr := netdetect.TunInterfaces(ctx, r)
	tunDetail := strings.Join(tuns, ", ")
	if tunErr != nil {
		tunDetail = errDetail(tunErr, tunDetail)
	}
	res.Checks = append(res.Checks, proto.Check{
		Name:   "tun",
		OK:     tunErr == nil && len(tuns) > 0,
		Warn:   true,
		Detail: tunDetail,
	})

	// A corrupt state file must be reported rather than silently reported as
	// "disabled".
	st, loadErr := state.Load(plan.StatePath)
	if loadErr != nil {
		res.Checks = append(res.Checks, proto.Check{
			Name:   "state",
			OK:     false,
			Warn:   true,
			Detail: loadErr.Error(),
		})
	}
	res.State = &proto.StateInfo{
		Enabled:      st.Enabled,
		WANInterface: st.WANInterface,
		Gateway4:     st.Gateway4,
		Gateway6:     st.Gateway6,
		Mark:         st.Mark,
	}
	return res
}

func engineAction(ctx context.Context, r execx.Runner, plan proto.Plan) proto.Result {
	res := proto.Result{Action: proto.ActionEngine}
	if plan.ServiceName == "" || !config.ValidServiceName(plan.ServiceName) {
		res.Error = fmt.Sprintf("invalid service name %q", plan.ServiceName)
		return res
	}
	switch plan.EngineOp {
	case proto.EngineStart:
		if err := engine.Start(ctx, r, plan.ServiceName); err != nil {
			res.Error = err.Error()
			return res
		}
	case proto.EngineStop:
		if err := engine.Stop(ctx, r, plan.ServiceName); err != nil {
			res.Error = err.Error()
			return res
		}
	case proto.EngineReload:
		if err := engine.Reload(ctx, r, plan.ServiceName); err != nil {
			res.Error = err.Error()
			return res
		}
	case proto.EngineStatus:
		active, err := engine.IsActive(ctx, r, plan.ServiceName)
		if err != nil {
			res.Error = err.Error()
			return res
		}
		res.OK = true
		res.Messages = append(res.Messages, fmt.Sprintf("%s active=%v", plan.ServiceName, active))
		return res
	default:
		res.Error = fmt.Sprintf("unknown engine operation %q", plan.EngineOp)
		return res
	}
	res.OK = true
	res.Changed = true
	return res
}

func initEngine(plan proto.Plan) proto.Result {
	res := proto.Result{Action: proto.ActionInitEngine}
	if plan.EngineConfig == "" || !filepath.IsAbs(plan.EngineConfig) {
		res.Error = "engine config path must be absolute"
		return res
	}
	if len(plan.EngineConfigBody) == 0 {
		res.Error = "no engine config body supplied"
		return res
	}
	// Reject anything that is not valid JSON before writing it to disk.
	var probe json.RawMessage
	if err := json.Unmarshal(plan.EngineConfigBody, &probe); err != nil {
		res.Error = "engine config is not valid JSON: " + err.Error()
		return res
	}
	if err := fsx.WriteFileAtomic(plan.EngineConfig, plan.EngineConfigBody, 0o644); err != nil {
		res.Error = err.Error()
		return res
	}
	res.OK = true
	res.Changed = true
	res.Messages = append(res.Messages, "wrote engine config to "+plan.EngineConfig)
	return res
}

func runChecks(ctx context.Context, r execx.Runner, plan proto.Plan, wan netdetect.WAN) ([]proto.Check, bool) {
	var checks []proto.Check
	fatalOK := true

	addFatal := func(name string, ok bool, detail string) {
		checks = append(checks, proto.Check{Name: name, OK: ok, Detail: detail, Fatal: true})
		if !ok {
			fatalOK = false
		}
	}

	exists, err := nft.Exists(ctx, r)
	if err != nil {
		addFatal("nft-table", false, err.Error())
	} else {
		addFatal("nft-table", exists, nft.Table())
	}

	for _, target := range []struct {
		family   int
		gateway  string
		table    int
		priority int
	}{
		{family: 4, gateway: wan.Gateway4, table: plan.TableV4, priority: plan.PriorityV4},
		{family: 6, gateway: wan.Gateway6, table: plan.TableV6, priority: plan.PriorityV6},
	} {
		if target.gateway == "" {
			continue
		}
		ruleOK, err := iprule.RuleExists(ctx, r, iprule.Spec{Family: target.family, Mark: plan.Mark, Table: target.table, Priority: target.priority})
		addFatal(fmt.Sprintf("ipv%d-rule", target.family), ruleOK && err == nil, errDetail(err, "fwmark rule present"))
		routeOK, err := iprule.RouteExists(ctx, r, target.family, target.table, target.gateway, wan.Interface)
		addFatal(fmt.Sprintf("ipv%d-route", target.family), routeOK && err == nil, errDetail(err, "table "+strconv.Itoa(target.table)))
	}

	if plan.ManageEngine && !plan.SkipEngine && plan.ServiceName != "" {
		active, err := engine.IsActive(ctx, r, plan.ServiceName)
		addFatal("engine-active", active && err == nil, errDetail(err, plan.ServiceName))
		tuns, tunErr := netdetect.TunInterfaces(ctx, r)
		detail := strings.Join(tuns, ", ")
		if tunErr != nil {
			detail = errDetail(tunErr, detail)
		}
		checks = append(checks, proto.Check{Name: "tun", OK: tunErr == nil && len(tuns) > 0, Warn: true, Detail: detail})
	}

	if plan.HealthProbe != "" {
		if err := health.DialTCP(ctx, plan.HealthProbe, healthTimeout); err != nil {
			checks = append(checks, proto.Check{
				Name:   "connectivity",
				OK:     false,
				Warn:   true,
				Detail: fmt.Sprintf("%v (set health.probe, or \"\" to disable)", err),
			})
		} else {
			checks = append(checks, proto.Check{Name: "connectivity", OK: true, Detail: "tcp " + plan.HealthProbe})
		}
	}

	return checks, fatalOK
}

func errDetail(err error, ok string) string {
	if err != nil {
		return err.Error()
	}
	return ok
}

// packageHints names the package that provides each required binary, so the
// preflight failure is actionable on any distribution.
var packageHints = map[string]string{
	"nft":        "install nftables",
	"ip":         "install iproute2",
	"systemctl":  "install systemd",
	"journalctl": "install systemd",
}

// preflight verifies the external commands required for an operation exist.
func preflight(r execx.Runner, needSystemctl bool) error {
	bins := []string{"nft", "ip"}
	if needSystemctl {
		bins = append(bins, "systemctl")
	}
	for _, bin := range bins {
		if _, err := r.LookPath(bin); err != nil {
			return fmt.Errorf("required binary %q not found in PATH; %s to continue", bin, packageHints[bin])
		}
	}
	return nil
}

// validatePlan re-validates every field of an untrusted plan before use.
func validatePlan(plan proto.Plan) error {
	if plan.Mark == 0 {
		return errors.New("mark must be non-zero")
	}
	if plan.WANInterface != "" && !config.ValidInterfaceName(plan.WANInterface) {
		return fmt.Errorf("invalid wan_interface %q", plan.WANInterface)
	}
	if plan.WANGateway4 != "" && net.ParseIP(plan.WANGateway4) == nil {
		return fmt.Errorf("invalid wan_gateway_v4 %q", plan.WANGateway4)
	}
	if plan.WANGateway6 != "" && net.ParseIP(plan.WANGateway6) == nil {
		return fmt.Errorf("invalid wan_gateway_v6 %q", plan.WANGateway6)
	}
	for _, check := range []struct {
		field string
		value int
	}{
		{"table_v4", plan.TableV4},
		{"table_v6", plan.TableV6},
		{"bypass_table_v4", plan.BypassTableV4},
		{"bypass_table_v6", plan.BypassTableV6},
	} {
		if err := config.CheckTableID(check.field, check.value); err != nil {
			return err
		}
	}
	for _, check := range []struct {
		field string
		value int
	}{
		{"priority_v4", plan.PriorityV4},
		{"priority_v6", plan.PriorityV6},
		{"bypass_priority_v4", plan.BypassPrioV4},
		{"bypass_priority_v6", plan.BypassPrioV6},
	} {
		if err := config.CheckPriority(check.field, check.value); err != nil {
			return err
		}
	}
	if plan.ManageEngine && plan.ServiceName != "" && !config.ValidServiceName(plan.ServiceName) {
		return fmt.Errorf("invalid service name %q", plan.ServiceName)
	}
	if plan.EngineConfig != "" && !filepath.IsAbs(plan.EngineConfig) {
		return fmt.Errorf("engine_config %q must be absolute", plan.EngineConfig)
	}
	if plan.StatePath == "" || !filepath.IsAbs(plan.StatePath) {
		return fmt.Errorf("state_path %q must be an absolute path", plan.StatePath)
	}
	for _, cidr := range plan.BypassCIDRs {
		if _, _, err := net.ParseCIDR(cidr); err != nil {
			return fmt.Errorf("invalid bypass CIDR %q", cidr)
		}
	}
	return nil
}