// SPDX-License-Identifier: MIT

package cli

import (
	"context"
	"fmt"
	"os"
	"regexp"
	"strconv"

	"github.com/spf13/cobra"

	"prx/internal/client"
	"prx/internal/engine"
	"prx/internal/execx"
	"prx/internal/proto"
	"prx/internal/state"
)

var markRE = regexp.MustCompile(`"(?:fwmark|default_mark|mark)"\s*:\s*(0x[0-9a-fA-F]+|\d+)`)

// nftVersionRE extracts the version from output such as
// "nftables v1.0.9 (Sunday, 1 September 2024)".
var nftVersionRE = regexp.MustCompile(`v(\d+)\.(\d+)(?:\.(\d+))?`)

// nftVersionWarning returns a warning when the installed nftables predates the
// `destroy` keyword (added in 0.9.4), or an empty string when it is new enough
// or the version cannot be determined.
func nftVersionWarning(output string) string {
	m := nftVersionRE.FindStringSubmatch(output)
	if m == nil {
		return ""
	}
	major, _ := strconv.Atoi(m[1])
	minor, _ := strconv.Atoi(m[2])
	patch := 0
	if m[3] != "" {
		patch, _ = strconv.Atoi(m[3])
	}
	if major == 0 && (minor < 9 || (minor == 9 && patch < 4)) {
		return fmt.Sprintf("nftables %d.%d.%d is too old; prx needs >= 0.9.4 (for the 'destroy' keyword)", major, minor, patch)
	}
	return ""
}

// diagnostic pairs a check with whether it was produced locally (and is
// therefore authoritative even when prx is disabled).
type diagnostic struct {
	check proto.Check
	local bool
}

func (a *app) newDoctorCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "doctor",
		Short: "Diagnose configuration and applied rules",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			cfg, err := a.loadConfig()
			if err != nil {
				return err
			}
			ctx := cmd.Context()
			r := execx.OSRunner{}

			// Local checks are always authoritative, even when prx is off.
			var diags []diagnostic
			addLocal := func(c proto.Check) { diags = append(diags, diagnostic{check: c, local: true}) }

			for _, bin := range []string{"nft", "ip", "systemctl"} {
				_, err := r.LookPath(bin)
				addLocal(proto.Check{Name: "binary:" + bin, OK: err == nil, Fatal: true, Detail: ternary(err == nil, "found", "missing — "+packageHint(bin))})
			}
			if _, err := r.LookPath("journalctl"); err != nil {
				// Not fatal: only 'prx logs' needs it.
				addLocal(proto.Check{Name: "binary:journalctl", OK: false, Warn: true, Detail: "missing — 'prx logs' unavailable"})
			}
			if _, err := r.LookPath("nft"); err == nil {
				res, verr := r.Run(ctx, nil, "nft", "--version")
				switch {
				case verr != nil:
					addLocal(proto.Check{Name: "nft-version", OK: false, Warn: true, Detail: verr.Error()})
				default:
					if w := nftVersionWarning(string(res.Stdout)); w != "" {
						addLocal(proto.Check{Name: "nft-version", OK: false, Warn: true, Detail: w})
					}
				}
			}
			addLocal(a.privilegeCheck(ctx, r))
			addLocal(proto.Check{Name: "config", OK: true, Detail: a.configPath})

			if cfg.Engine.Manage && cfg.Engine.ConfigPath != "" {
				warnings, err := engine.ValidateConfig(cfg.Engine.ConfigPath)
				switch {
				case err != nil:
					addLocal(proto.Check{Name: "engine-config", OK: false, Warn: true, Detail: err.Error()})
				default:
					if len(warnings) == 0 {
						addLocal(proto.Check{Name: "engine-config", OK: true, Detail: cfg.Engine.ConfigPath})
					}
					for _, w := range warnings {
						addLocal(proto.Check{Name: "engine-config", OK: false, Warn: true, Detail: w})
					}
				}
				if collisions := markCollisions(cfg.Engine.ConfigPath, cfg.Marks.InboundMark); len(collisions) > 0 {
					addLocal(proto.Check{Name: "mark-collision", OK: false, Fatal: true,
						Detail: fmt.Sprintf("engine config reuses mark 0x%x, same as prx (%s)", cfg.Marks.InboundMark, collisions[0])})
				}
			}

			st, stateErr := state.Load(cfg.State.Path)
			if stateErr != nil {
				addLocal(proto.Check{Name: "state", OK: false, Warn: true, Detail: stateErr.Error()})
			} else {
				addLocal(proto.Check{Name: "state", OK: true, Detail: fmt.Sprintf("enabled=%v", st.Enabled)})
			}

			// Deep inspection needs root; degrade gracefully.
			res, err := client.Dispatch(ctx, a.buildPlan(cfg, proto.ActionInspect))
			if err != nil {
				diags = append(diags, diagnostic{check: proto.Check{Name: "inspect", OK: false, Warn: true, Detail: err.Error()}})
			} else {
				// A remote check must not repeat a local one (for example
				// "state"): the local result is authoritative.
				have := make(map[string]bool, len(diags))
				for _, d := range diags {
					have[d.check.Name] = true
				}
				for _, c := range res.Checks {
					if have[c.Name] {
						continue
					}
					have[c.Name] = true
					diags = append(diags, diagnostic{check: c})
				}
			}

			fmt.Fprintln(a.out, "prx doctor")
			failed := false
			for _, d := range diags {
				c := d.check
				mark := checkMark(c)
				if c.Fatal && !c.OK && !c.Warn && !d.local && !st.Enabled {
					// Expected when prx is disabled.
					mark = "off"
				}
				if c.Fatal && !c.OK && !c.Warn && (d.local || st.Enabled) {
					failed = true
				}
				line := fmt.Sprintf("  [%s] %s", mark, c.Name)
				if c.Detail != "" {
					line += ": " + c.Detail
				}
				fmt.Fprintln(a.out, line)
			}
			if failed {
				return fmt.Errorf("doctor found critical problems")
			}
			return nil
		},
	}
	return cmd
}

// packageHint names the package that provides a required binary, so a missing
// dependency is actionable on any distribution.
func packageHint(bin string) string {
	switch bin {
	case "nft":
		return "install nftables"
	case "ip":
		return "install iproute2"
	case "systemctl", "journalctl":
		return "install systemd"
	default:
		return "install it"
	}
}

// privilegeCheck reports how prx will reach root, or what prevents it from
// doing so.
func (a *app) privilegeCheck(_ context.Context, r execx.Runner) proto.Check {
	const name = "privilege"
	if os.Geteuid() == 0 {
		return proto.Check{Name: name, OK: true, Detail: "running as root (helper runs in-process)"}
	}
	helperPath, err := client.FindHelper()
	if err != nil {
		return proto.Check{Name: name, OK: false, Fatal: true, Detail: err.Error()}
	}
	detail := "helper " + helperPath + ", via "
	if _, err := r.LookPath("pkexec"); err == nil {
		return proto.Check{Name: name, OK: true, Detail: detail + "pkexec (graphical prompt)"}
	}
	if _, err := r.LookPath("sudo"); err == nil {
		return proto.Check{Name: name, OK: true, Detail: detail + "sudo (passwordless required; or run 'sudo prx ...')"}
	}
	return proto.Check{
		Name:   name,
		OK:     false,
		Fatal:  true,
		Detail: "neither pkexec nor sudo found; install polkit or sudo, or run prx as root",
	}
}

func markCollisions(path string, mark uint32) []string {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil
	}
	var out []string
	for _, m := range markRE.FindAllStringSubmatch(string(raw), -1) {
		v, err := strconv.ParseUint(m[1], 0, 32)
		if err != nil {
			continue
		}
		if uint32(v) == mark {
			out = append(out, m[0])
		}
	}
	return out
}

func ternary(cond bool, a, b string) string {
	if cond {
		return a
	}
	return b
}