package cli

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"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
		model   string
	)

	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" +
			"For llama.cpp on a machine with NVIDIA GPUs it then loads a small GGUF\n" +
			"(--model, or the smallest in the weights cache) with every layer on the GPU\n" +
			"and asks for one token: a build whose CUDA backend does not load computes on\n" +
			"the CPU, and that is where it is caught.\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{}
			// The build the agent installed, when none is named: the newest
			// under /opt/llmgw/engines, which is where every node keeps it.
			if bin == "" && name == "llamacpp" {
				bin = installedLlama()
			}
			if bin != "" {
				switch name {
				case "llamacpp":
					cfg.LlamaCpp.Bin = bin
				case "vllm":
					cfg.VLLM.Bin = bin
				case "sglang":
					cfg.SGLang.Python = bin
				case "trtllm":
					cfg.TRTLLM.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(), eng, 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)
				// Seeing a GPU is not computing on it: a small model loaded
				// on it and asked for a token is.
				if b, ok := eng.(interface{ Binary() string }); ok && deviceErr == nil && model != "none" {
					deviceErr = computeOnGPU(cmd.Context(), b.Binary(), model, 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 (default: the build the agent installed, else PATH)")
	f.StringVar(&model, "model", "", "a GGUF file to load on the GPU and ask for a token (default: the smallest "+
		"in this node's weights cache; none skips it)")
	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. For an engine that runs a replica on one node,
// that is one node.
func representativeShard(engine string, nodes, tp, perNode int) engines.ShardSpec {
	if nodes < 1 || engine == "trtllm" {
		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, eng engines.Engine, c engines.Command) (string, error) {
	// Everything up to the first flag, plus the help request. Passing the real
	// arguments would start the engine.
	path := c.Path
	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...)
	// An engine started through a launcher of llmgw's is asked through its
	// own entry point: the interpreter's --help describes the interpreter.
	if h, ok := eng.(engines.HelpCommand); ok {
		path, args = h.HelpCommand(c)
	}

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

	ec := exec.CommandContext(ctx, 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)",
			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
}

// installedLlama is the newest llama-server the agent installed here.
func installedLlama() string {
	m, _ := filepath.Glob("/opt/llmgw/engines/llamacpp/*/llama-server")
	if len(m) == 0 {
		return ""
	}
	sort.Strings(m)
	return m[len(m)-1]
}

// weightDirs are where a node keeps models it has fetched.
var weightDirs = []string{"/var/lib/llmgw-weights", "/var/lib/llmgw-models", "/srv/models"}

// smallestGGUF is the smallest GGUF file under dirs, for a check that only
// needs something to load.
func smallestGGUF(dirs []string) string {
	best, size := "", int64(0)
	for _, d := range dirs {
		_ = filepath.WalkDir(d, func(p string, e os.DirEntry, err error) error {
			if err != nil || e.IsDir() || !strings.HasSuffix(strings.ToLower(p), ".gguf") ||
				strings.Contains(strings.ToLower(filepath.Base(p)), "mmproj") {
				return nil
			}
			if info, err := e.Info(); err == nil && (best == "" || info.Size() < size) {
				best, size = p, info.Size()
			}
			return nil
		})
	}
	return best
}

// computeOnGPU loads a model with every layer on the GPU, asks for one
// token, and reads from llama.cpp's own output where the layers went.
func computeOnGPU(ctx context.Context, bin, model string, out io.Writer) error {
	if _, err := exec.LookPath("nvidia-smi"); err != nil {
		return nil
	}
	if model == "" {
		if model = smallestGGUF(weightDirs); model == "" {
			fmt.Fprintln(out, "no GGUF in this node's weights cache to compute with; --model PATH runs one")
			return nil
		}
	}
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return err
	}
	port := ln.Addr().(*net.TCPAddr).Port
	_ = ln.Close()
	ctx, cancel := context.WithTimeout(ctx, 3*time.Minute)
	defer cancel()
	var logs syncBuffer
	c := exec.CommandContext(ctx, bin, "-m", model, "--host", "127.0.0.1", "--port", strconv.Itoa(port), //nolint:gosec // the engine's own binary
		"--n-gpu-layers", "999", "--ctx-size", "512", "--parallel", "1",
		// Where the layers went is said only at this verbosity.
		"-lv", "4")
	c.Stdout, c.Stderr = &logs, &logs
	if err := c.Start(); err != nil {
		return fmt.Errorf("starting %s: %w", bin, err)
	}
	defer func() { _ = c.Process.Kill(); _ = c.Wait() }()
	base := "http://127.0.0.1:" + strconv.Itoa(port)
	start := time.Now()
	for {
		resp, err := http.Get(base + "/health") //nolint:gosec,noctx // loopback, bounded by the loop
		if err == nil {
			_ = resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				break
			}
		}
		if ctx.Err() != nil {
			return fmt.Errorf("%s did not come up with %s: %s", bin, filepath.Base(model), tailLines(logs.String(), 5))
		}
		time.Sleep(500 * time.Millisecond)
	}
	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, base+"/completion",
		strings.NewReader(`{"prompt":"Hello","n_predict":1}`))
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("asking it for a token: %w", err)
	}
	_ = resp.Body.Close()
	took := time.Since(start)
	comp, said := engines.ReadLlamaCompute(strings.Split(logs.String(), "\n"))
	if !said {
		fmt.Fprintf(out, "computed a token with %s, but llama.cpp did not say where\n", filepath.Base(model))
		return nil
	}
	if comp.OnTheCPU() {
		return fmt.Errorf("%s computed on the CPU though every layer was asked for on the GPU (%d of %d offloaded): "+
			"its CUDA backend did not load; ldd on its libggml-cuda.so names the missing library",
			filepath.Base(model), comp.Offloaded, comp.Total)
	}
	fmt.Fprintf(out, "computed a token on %s: %s, %d of %d layers offloaded, in %s\n", comp.Backend,
		strings.Join(comp.Devices, "; "), comp.Offloaded, comp.Total, took.Round(100*time.Millisecond))
	return nil
}

func tailLines(s string, n int) string {
	l := strings.Split(strings.TrimSpace(s), "\n")
	if len(l) > n {
		l = l[len(l)-n:]
	}
	return strings.Join(l, " | ")
}

// syncBuffer is a buffer two writers share.
type syncBuffer struct {
	mu sync.Mutex
	b  strings.Builder
}

func (s *syncBuffer) Write(p []byte) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.b.Write(p)
}

func (s *syncBuffer) String() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.b.String()
}