// Package engines builds the command lines and environments that start
// inference engines (plan §5.3).
//
// Nothing here starts a process; that is the supervisor's job. Keeping command
// construction pure makes it testable without GPUs, which matters because the
// environment generation is where multi-node runs go wrong: a single missing
// NCCL variable turns a working tensor-parallel group into one that hangs on
// the first all-reduce.
package engines

import (
	"context"
	"fmt"
	"os"
	"slices"
	"sort"
	"strconv"
	"strings"
)

// Profile is how a model runs on a GPU class. It mirrors the catalog profile
// (plan §6.1), reduced to what launching needs.
type Profile struct {
	Name     string
	Engine   string
	GPUClass string

	// TP and PP are tensor and pipeline parallel degrees.
	TP int
	PP int
	// GPUsPerNode is how many GPUs each shard occupies.
	GPUsPerNode int
	// Nodes is how many nodes the replica spans. Above 1, a fabric is required.
	Nodes int

	// Args are passed through to the engine untouched.
	Args []string
	// Env overrides anything the engine adapter generates.
	Env map[string]string

	ModelPath string
	// ServedModelName is the name clients use, which need not match the path.
	ServedModelName string
	// Adapters are LoRA adapters to load beside the model, in order: an
	// adapter's index is the id requests apply it by.
	Adapters []Adapter
	// KVEvents has the engine publish its KV cache events to the agent,
	// for the router's precise index. vLLM only.
	KVEvents bool `json:"kv_events,omitempty"`
}

// ShardSpec is one shard of one replica on this node.
type ShardSpec struct {
	ReplicaID string
	ShardID   string
	// Rank is the node rank; 0 is the leader and serves HTTP.
	Rank   int
	NNodes int

	// LeaderAddr and LeaderPort are the distributed init endpoint, always on the
	// management network rather than the fabric: NCCL bootstraps over TCP and
	// only then moves to RDMA.
	LeaderAddr string
	LeaderPort int

	// GPUIndices are the local GPUs this shard uses.
	GPUIndices []int
	// GPUUUIDs name the same devices by UUID, in the same order, and are set
	// by the agent only when one of them is a MIG instance. An instance has
	// no index CUDA knows: CUDA_VISIBLE_DEVICES=MIG-<uuid> is the only way to
	// give a process one, and the indices placement uses for it (1000 and
	// up) mean nothing to CUDA.
	GPUUUIDs []string `json:"gpu_uuids,omitempty"`
	// HCAs are the adapters of the chosen fabric, and only that fabric. A shard
	// must never be given HCAs from two fabrics.
	HCAs []string
	// MgmtInterface is the ethernet interface used for bootstrap traffic.
	MgmtInterface string
	// GIDIndex is the RoCE v2 GID index for the HCAs, set by the agent when
	// they are RoCE; nil on InfiniBand, where NCCL needs none.
	GIDIndex *int

	HTTPPort int
	Profile  Profile
	// BindAddr is the address the engine listens on: the node's management
	// address, which is what the router reaches it on, set by the agent.
	// Empty listens on every address.
	BindAddr string `json:"bind_addr,omitempty"`

	// APIKey is what the engine demands of every request. Engines listen on
	// the node's management address; without a key anybody who can reach it
	// talks to the model directly, past the gateway's authentication, quotas
	// and accounting. Given to the engine in its environment rather than on
	// its command line, which every user on the node can read.
	APIKey string `json:"api_key,omitempty"`
	// KVEventsPort is the loopback port the agent receives this shard's KV
	// cache events on, chosen by the agent; zero when it publishes none.
	// Kept with the spec so an agent that restarts listens where the engine
	// it adopts is still sending.
	KVEventsPort int `json:"kv_events_port,omitempty"`
	// Threads is the CPU cores placement reserved for a shard with no GPU,
	// and so what its engine computes on. Zero leaves it to the engine's own
	// configuration.
	Threads int `json:"threads,omitempty"`

	// Pool and Fabric are control plane concepts the agent cannot derive. They
	// are carried so the agent can report them back after a restart; Fabric
	// also says, with an "nvl:" prefix, that the replica is inside a
	// multi-node NVLink domain, which changes NCCL's environment.
	Pool   string
	Fabric string
	// ProfileFP identifies the version of the profile this shard was started
	// from. Carried for the same reason as Pool and Fabric, and never read: it
	// lets the control plane tell, after its own restart, whether a shard that
	// is still running matches the profile the catalog holds now.
	ProfileFP string
}

// Host is where the engine listens.
func (s ShardSpec) Host() string {
	if s.BindAddr == "" {
		return "0.0.0.0"
	}
	return s.BindAddr
}

// Leader reports whether this shard serves HTTP for the replica.
func (s ShardSpec) Leader() bool { return s.Rank == 0 }

// Command is a fully resolved launch: what to exec, with which environment.
type Command struct {
	Path string
	Args []string
	// Env is the complete environment, not a delta.
	Env []string
	// Name distinguishes the processes of a multi-process shard in logs.
	Name string
	// WaitForLeaderPort makes the supervisor hold the next command until
	// LeaderAddr:LeaderPort accepts a connection.
	//
	// It exists for Ray: the head has to be listening before a worker can join,
	// and both sides express the same wait because both dial the same address.
	WaitForLeaderPort bool
}

// Engine builds launch commands for one inference engine.
type Engine interface {
	// Name is the engine's identifier in the catalog.
	Name() string
	// Validate rejects a profile that cannot work on these GPUs, before
	// anything is started.
	Validate(p Profile, gpus []GPUInfo) error

	// Commands returns the processes that make up one shard, in start order.
	//
	// Most shards are a single process. vLLM under Ray is not: the leader runs
	// a Ray head and then the server against it, and a worker runs only the Ray
	// join. Returning a sequence rather than one command is what lets the
	// supervisor treat both the same way — every command but the last is a
	// sidecar kept alive for the shard's lifetime, and the last one's exit is
	// the shard's exit.
	Commands(s ShardSpec) ([]Command, error)

	// ServesHTTP reports whether this rank answers the OpenAI API.
	//
	// In a multi-node replica only the leader does. Health-checking a worker's
	// port would time out and fail a replica that is coming up perfectly well,
	// which is exactly what happened before this existed.
	ServesHTTP(s ShardSpec) bool

	// HealthPath and MetricsPath are appended to the shard's base URL.
	HealthPath() string
	MetricsPath() string
}

// CPUFallback is an engine that can come up on the CPU when it was placed on
// GPUs, and keep serving there without an error. The agent checks such a shard
// is on its GPUs before calling it ready.
type CPUFallback interface {
	FallsBackToCPU() bool
	// Devices lists the GPUs the engine's build can see, for saying why.
	Devices(ctx context.Context) ([]string, error)
}

// HelpCommand is an engine whose program is not asked for its help by
// running the command's path with --help, because the path is an interpreter
// starting it through a launcher.
type HelpCommand interface {
	HelpCommand(c Command) (path string, args []string)
}

// GPUInfo is the subset of GPU facts validation needs.
type GPUInfo struct {
	Index             int
	Name              string
	ComputeCapability string
	MemoryTotalMB     int
}

// Registry holds the known engines.
type Registry struct {
	engines map[string]Engine
}

// NewRegistry returns a registry with the built-in engines.
func NewRegistry(cfg Config) *Registry {
	r := &Registry{engines: map[string]Engine{}}
	r.Add(&VLLM{cfg: cfg.VLLM})
	r.Add(&SGLang{cfg: cfg.SGLang})
	r.Add(&LlamaCpp{cfg: cfg.LlamaCpp})
	r.Add(&TRTLLM{cfg: cfg.TRTLLM})
	r.Add(&Sim{cfg: cfg.Sim})
	return r
}

// UseInstalled points an engine at a build the agent installed, and reports
// whether the engine takes one.
func (r *Registry) UseInstalled(name, path string) bool {
	u, ok := r.engines[name].(interface{ UseInstalled(string) })
	if ok {
		u.UseInstalled(path)
	}
	return ok
}

// Add registers an engine, replacing any engine of the same name.
func (r *Registry) Add(e Engine) { r.engines[e.Name()] = e }

// Get returns an engine by name.
func (r *Registry) Get(name string) (Engine, error) {
	e, ok := r.engines[name]
	if !ok {
		return nil, fmt.Errorf("unknown engine %q (have %s)", name, strings.Join(r.Names(), ", "))
	}
	return e, nil
}

// Names lists the registered engines.
func (r *Registry) Names() []string {
	out := make([]string, 0, len(r.engines))
	for n := range r.engines {
		out = append(out, n)
	}
	slices.Sort(out)
	return out
}

// Config configures the engine adapters.
type Config struct {
	// Disabled are engines this node never installs or advertises, whatever
	// the control plane asks: vLLM on a node of cards no build of it runs on,
	// that is meant for llama.cpp alone.
	Disabled []string       `koanf:"disabled"`
	VLLM     VLLMConfig     `koanf:"vllm"`
	SGLang   SGLangConfig   `koanf:"sglang"`
	LlamaCpp LlamaCppConfig `koanf:"llamacpp"`
	// TRTLLM is TensorRT-LLM, which the agent does not install: a node
	// offers it only when its interpreter is named here.
	TRTLLM TRTLLMConfig `koanf:"trtllm"`
	Sim    SimConfig    `koanf:"sim"`
}

// --- shared helpers ---

// NCCLEnv is the environment a shard with this spec is given, for a tool that
// shows an operator what an engine would see.
func NCCLEnv(s ShardSpec) map[string]string { return baseEnv(s) }

// baseEnv builds the environment every engine shares.
//
// This is the part of §5.3 the plan marks critical. The rules encoded here:
//
//   - CUDA_VISIBLE_DEVICES pins the shard to its GPUs, so two shards on one
//     node cannot fight over the same card.
//   - NCCL_IB_HCA lists only the chosen fabric's adapters. Leaving it unset
//     lets NCCL pick adapters from another fabric, where the peer is
//     unreachable, and the job hangs instead of failing.
//   - NCCL_SOCKET_IFNAME pins bootstrap to the management interface. NCCL
//     otherwise probes every interface, including IB-over-IP addresses that
//     cannot route between fabrics.
//   - NCCL_IB_DISABLE is set when there is no fabric, so NCCL stops looking for
//     one instead of timing out.
func baseEnv(s ShardSpec) map[string]string {
	env := map[string]string{
		"CUDA_VISIBLE_DEVICES": visibleDevices(s),
		// The indices above are nvidia-smi's, which counts by PCI bus. CUDA
		// counts fastest first unless told otherwise, and on a node with two
		// kinds of GPU the shard would get a different card than placement
		// gave it — one another shard may hold.
		"CUDA_DEVICE_ORDER": "PCI_BUS_ID",
	}

	if len(s.HCAs) > 0 {
		env["NCCL_IB_HCA"] = strings.Join(s.HCAs, ",")
		env["NCCL_IB_DISABLE"] = "0"
		// RoCE: the index of the routable RoCE v2 entry. NCCL's default is a
		// v1 link-local one on most hosts, and a job on it hangs.
		if s.GIDIndex != nil {
			env["NCCL_IB_GID_INDEX"] = strconv.Itoa(*s.GIDIndex)
		}
	} else {
		// Single node, or a node with no fabric. Either way IB is not the path.
		env["NCCL_IB_DISABLE"] = "1"
	}

	// Placed inside a multi-node NVLink domain: NCCL is to carry the
	// replica's collectives over NVLink between its nodes, which it only
	// does when told to, and needs IMEX for.
	if s.NNodes > 1 && strings.HasPrefix(s.Fabric, "nvl:") {
		env["NCCL_MNNVL_ENABLE"] = "1"
		env["NCCL_CUMEM_ENABLE"] = "1"
	}

	if s.MgmtInterface != "" {
		env["NCCL_SOCKET_IFNAME"] = s.MgmtInterface
		env["GLOO_SOCKET_IFNAME"] = s.MgmtInterface
	}

	// Profile overrides win: they are the escape hatch for the per-cluster
	// tuning no adapter can guess.
	for k, v := range s.Profile.Env {
		// Checked again here, not only where the model was defined: this is
		// the process about to run as root, whatever sent the profile.
		if bad, _ := UnsafeEnv(k); bad {
			continue
		}
		env[k] = v
	}
	return env
}

// materializeEnv merges generated variables over the current process
// environment and returns it in exec form.
//
// Inheriting matters: engines need PATH, HOME, HF_HOME, LD_LIBRARY_PATH and
// whatever else the operator set up in the unit file.
func materializeEnv(generated map[string]string) []string {
	merged := map[string]string{}
	for _, kv := range os.Environ() {
		if k, v, ok := strings.Cut(kv, "="); ok {
			merged[k] = v
		}
	}
	for k, v := range generated {
		merged[k] = v
	}

	out := make([]string, 0, len(merged))
	for k, v := range merged {
		out = append(out, k+"="+v)
	}
	// Sorted so a command is reproducible and diffable in logs.
	sort.Strings(out)
	return out
}

// visibleDevices is CUDA_VISIBLE_DEVICES for a shard: the indices placement
// gave it, or their UUIDs when the agent named them, which it does for a MIG
// instance.
func visibleDevices(s ShardSpec) string {
	if len(s.GPUUUIDs) > 0 {
		return strings.Join(s.GPUUUIDs, ",")
	}
	return joinInts(s.GPUIndices)
}

// OnMIG reports whether the shard was given a MIG instance.
func (s ShardSpec) OnMIG() bool {
	return slices.ContainsFunc(s.GPUUUIDs, func(u string) bool { return strings.HasPrefix(u, "MIG-") })
}

func joinInts(v []int) string {
	parts := make([]string, len(v))
	for i, n := range v {
		parts[i] = strconv.Itoa(n)
	}
	return strings.Join(parts, ",")
}

// validateCommon holds the checks that apply to every engine.
func validateCommon(p Profile, gpus []GPUInfo) error {
	if p.ModelPath == "" {
		return fmt.Errorf("profile %q: model path is empty", p.Name)
	}
	if p.TP <= 0 {
		return fmt.Errorf("profile %q: tensor parallel degree must be positive", p.Name)
	}
	if p.Nodes > 1 && p.GPUsPerNode <= 0 {
		return fmt.Errorf("profile %q: gpus_per_node is required for a %d-node profile", p.Name, p.Nodes)
	}

	perNode := p.GPUsPerNode
	if perNode == 0 {
		perNode = p.TP
	}
	nodes := max(p.Nodes, 1)
	if want := perNode * nodes; want != p.TP*max(p.PP, 1) {
		return fmt.Errorf("profile %q: %d nodes x %d GPUs does not match TP %d x PP %d",
			p.Name, nodes, perNode, p.TP, max(p.PP, 1))
	}
	if len(gpus) > 0 && len(gpus) < perNode {
		return fmt.Errorf("profile %q: needs %d GPUs per node, node has %d", p.Name, perNode, len(gpus))
	}

	// A tensor-parallel group must be uniform: mixing classes makes the slowest
	// card set the pace for every collective, and mismatched memory makes the
	// smallest card decide what fits.
	if len(gpus) > 1 {
		first := gpus[0].Name
		for _, g := range gpus[1:] {
			if g.Name != first {
				return fmt.Errorf("profile %q: GPUs are not uniform (%q and %q); "+
					"a tensor-parallel group may not mix GPU types", p.Name, first, g.Name)
			}
		}
	}
	return nil
}

// requireComputeCapability rejects profiles whose quantization the hardware
// cannot execute.
func requireComputeCapability(p Profile, gpus []GPUInfo, minCap float64, why string) error {
	for _, g := range gpus {
		cap, err := strconv.ParseFloat(g.ComputeCapability, 64)
		if err != nil {
			// Unknown capability: let it through rather than block a launch on
			// a parsing problem. The engine will refuse it clearly enough.
			continue
		}
		if cap < minCap {
			return fmt.Errorf("profile %q: %s needs compute capability %.1f, GPU %d (%s) has %.1f",
				p.Name, why, minCap, g.Index, g.Name, cap)
		}
	}
	return nil
}

// quantizationOf extracts the --quantization argument, if present.
func quantizationOf(args []string) string {
	for i, a := range args {
		if a == "--quantization" && i+1 < len(args) {
			return args[i+1]
		}
		if v, ok := strings.CutPrefix(a, "--quantization="); ok {
			return v
		}
	}
	return ""
}

// Adapter is a LoRA adapter an engine loads beside its model. Path is a
// source until the agent has fetched it, a local file after.
type Adapter struct {
	Name string `json:"name"`
	Path string `json:"path"`
}