package cli

import (
	"context"
	"fmt"
	"io"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
)

func newEngineCheckCmd() *cobra.Command {
	var (
		help    string
		bin     string
		nodes   int
		tp      int
		perNode int
	)

	cmd := &cobra.Command{
		Use:   "engine-check <engine>",
		Short: "Check that an engine still accepts the flags llmgw builds for it",
		Long: "llmgw's launch arguments are read from each engine's documentation, and\n" +
			"engines rename and remove flags between versions. When that happens the\n" +
			"failure arrives minutes into a model load, in an engine's own error\n" +
			"message, on a node nobody was watching.\n\n" +
			"This asks the engine itself. It runs the engine's --help and checks every\n" +
			"flag llmgw would pass against it, which needs no GPU — a CLI prints its\n" +
			"help on any machine. Run it after upgrading an engine.\n\n" +
			"With --help-output it reads a saved --help instead of running anything,\n" +
			"so the check works from a laptop against help captured on a node.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			name := args[0]
			reg := engines.NewRegistry(engines.Config{})
			eng, err := reg.Get(name)
			if err != nil {
				return err
			}

			// The binary lives where the agent's configuration says, which is
			// rarely on PATH: a check that cannot find the engine reports
			// nothing rather than reporting a problem.
			cfg := engines.Config{}
			if bin != "" {
				switch name {
				case "llamacpp":
					cfg.LlamaCpp.Bin = bin
				case "vllm":
					cfg.VLLM.Bin = bin
				case "sglang":
					cfg.SGLang.Python = bin
				case "sim":
					cfg.Sim.Bin = bin
				}
				reg = engines.NewRegistry(cfg)
				if eng, err = reg.Get(name); err != nil {
					return err
				}
			}

			spec := representativeShard(name, nodes, tp, perNode)
			cmds, err := eng.Commands(spec)
			if err != nil {
				return fmt.Errorf("building a command for %s: %w", name, err)
			}

			out, errOut := cmd.OutOrStdout(), cmd.ErrOrStderr()
			var checked, skipped int
			var unknown []string

			// Per command, and that distinction is the whole point. A shard can
			// be several processes — vLLM under Ray is `ray start` and then
			// `vllm serve` — and they are different programs with different
			// vocabularies. Checking the sidecar's flags against the engine's
			// help reports --head and --num-gpus as missing, which is true of
			// vllm and irrelevant, because ray is what receives them.
			for i, c := range cmds {
				fmt.Fprintf(errOut, "  %s %s\n", c.Path, strings.Join(c.Args, " "))

				text := ""
				switch {
				case help != "" && i == len(cmds)-1:
					// A supplied help belongs to the engine, which is the last
					// command; the sidecars have their own and we do not have it.
					text = help
				case help != "":
					skipped++
					fmt.Fprintf(errOut, "    skipped: --help-output describes the engine, not %s\n", c.Name)
					continue
				default:
					var err error
					if text, err = runHelp(cmd.Context(), c); err != nil {
						skipped++
						fmt.Fprintf(errOut, "    skipped: %v\n", err)
						continue
					}
				}

				// A truncated help produces confident nonsense: every flag below
				// the cut looks missing. Real help documents --help itself,
				// which is the cheapest signal that we have all of it.
				if !mentions(text, "--help") {
					fmt.Fprintf(errOut, "    warning: this help does not document --help, "+
						"which usually means it was truncated; flags below the cut will look missing\n")
				}

				seen := map[string]bool{}
				for _, a := range c.Args {
					flag, _, _ := strings.Cut(a, "=")
					if !strings.HasPrefix(flag, "--") || seen[flag] {
						continue
					}
					seen[flag] = true
					checked++
					if !mentions(text, flag) {
						unknown = append(unknown, c.Name+": "+flag)
					}
				}
			}
			sort.Strings(unknown)

			// Flags accepted is not the same as running where it is placed: a
			// GPU build whose backend does not load accepts every flag and
			// computes on the CPU. Asked of the engine here, not of a help.
			var deviceErr error
			if fb, ok := eng.(engines.CPUFallback); ok && help == "" {
				deviceErr = checkDevices(cmd.Context(), name, fb, out)
			}

			fmt.Fprintf(out, "%s: %d flags checked", name, checked)
			if skipped > 0 {
				fmt.Fprintf(out, ", %d commands skipped", skipped)
			}
			fmt.Fprintln(out)

			if len(unknown) == 0 {
				if checked == 0 {
					return fmt.Errorf("nothing was checked; install the engine here or pass --help-output")
				}
				fmt.Fprintln(out, "every flag llmgw builds is accepted")
				return deviceErr
			}
			fmt.Fprintf(out, "\n%d flags not mentioned by the program they are passed to:\n", len(unknown))
			for _, f := range unknown {
				fmt.Fprintf(out, "  %s\n", f)
			}
			return fmt.Errorf("%s: %d flags are not accepted", name, len(unknown))
		},
	}

	f := cmd.Flags()
	f.StringVar(&help, "help-output", "", "a saved --help to check against, instead of running the engine")
	f.StringVar(&bin, "bin", "", "where the engine lives, when it is not on PATH")
	f.IntVar(&nodes, "nodes", 2, "shape the command for a replica of this many nodes")
	f.IntVar(&tp, "tp", 8, "tensor parallel degree")
	f.IntVar(&perNode, "gpus-per-node", 4, "GPUs each node contributes")
	return cmd
}

// representativeShard is the widest command llmgw would ever build, because a
// flag that only appears in the multi-node case is exactly the one nobody
// notices has been renamed.
func representativeShard(engine string, nodes, tp, perNode int) engines.ShardSpec {
	if nodes < 1 {
		nodes = 1
	}
	return engines.ShardSpec{
		ReplicaID: "check", ShardID: "check-r0", Rank: 0, NNodes: nodes,
		LeaderAddr: "10.0.0.1", LeaderPort: 20001,
		GPUIndices: seq(perNode), HCAs: []string{"mlx5_0"},
		MgmtInterface: "eth0", HTTPPort: 20000,
		Profile: engines.Profile{
			Name: "check", Engine: engine, TP: tp, Nodes: nodes, GPUsPerNode: perNode,
			ModelPath: "/models/x", ServedModelName: "x",
		},
	}
}

func seq(n int) []int {
	out := make([]int, n)
	for i := range out {
		out[i] = i
	}
	return out
}

// helpArgs is how each program is asked for its *complete* help.
//
// Not every program gives it all up for --help. vLLM prints a summary listing
// section names and tells you to ask again per section; checking against that
// reports every flag as missing, confidently and wrongly. This is exactly what
// happened the first time this ran against a real vLLM.
var helpArgs = map[string][]string{
	"vllm": {"--help=all"},
}

// runHelp asks one program to describe itself.
func runHelp(ctx context.Context, c engines.Command) (string, error) {
	// Everything up to the first flag, plus the help request. Passing the real
	// arguments would start the engine.
	var args []string
	for _, a := range c.Args {
		if strings.HasPrefix(a, "-") {
			break
		}
		args = append(args, a)
	}
	ask, ok := helpArgs[filepath.Base(c.Path)]
	if !ok {
		ask = []string{"--help"}
	}
	args = append(args, ask...)

	ctx, cancel := context.WithTimeout(ctx, 2*time.Minute)
	defer cancel()

	ec := exec.CommandContext(ctx, c.Path, args...)
	// vLLM refuses to do anything at all on a machine with no accelerator,
	// including print its help, unless it is told what to assume. Asking an
	// engine about its flags should not require the hardware it runs on.
	ec.Env = append(ec.Environ(), "VLLM_TARGET_DEVICE=cpu")

	raw, err := ec.CombinedOutput()
	if err != nil && len(raw) == 0 {
		return "", fmt.Errorf("could not run %s: %w (install it here, or pass --help-output)",
			c.Path, err)
	}
	return string(raw), nil
}

// mentions reports whether the help text documents a flag.
//
// Matched on a word boundary, so --port does not satisfy --port-range and a
// missing flag cannot hide behind a longer one that happens to contain it.
func mentions(help, flag string) bool {
	for i := 0; ; {
		j := strings.Index(help[i:], flag)
		if j < 0 {
			return false
		}
		end := i + j + len(flag)
		if end == len(help) || !isFlagChar(help[end]) {
			return true
		}
		i = end
	}
}

func isFlagChar(b byte) bool {
	return b == '-' || b == '_' || (b >= 'a' && b <= 'z') ||
		(b >= 'A' && b <= 'Z') || (b >= '0' && b <= '9')
}

// checkDevices says which GPUs an engine that can fall back to the CPU sees on
// this machine, and fails when the machine has NVIDIA GPUs and it sees none.
func checkDevices(ctx context.Context, name string, fb engines.CPUFallback, out io.Writer) error {
	if _, err := exec.LookPath("nvidia-smi"); err != nil {
		fmt.Fprintln(out, "no NVIDIA GPU here: which devices it computes on was not checked")
		return nil
	}
	devs, err := fb.Devices(ctx)
	if err != nil {
		return err
	}
	if len(devs) == 0 {
		return fmt.Errorf("this machine has NVIDIA GPUs and %s sees none: it would serve a GPU placement on the CPU. "+
			"Its GPU backend did not load; ldd on its libggml-cuda.so names the missing library", name)
	}
	fmt.Fprintf(out, "computes on: %s\n", strings.Join(devs, "; "))
	return nil
}