package cli

import (
	"fmt"
	"net/url"
	"strconv"
	"strings"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelserve"
	"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
)

// newModelPlanCmd answers whether a model fits this cluster, and how.
//
// The arithmetic it replaces is the last thing standing between having GPUs and
// serving a model on them. Deciding tp, gpu_class and a context that will
// actually fit means knowing the model's layer count, its KV head count, what
// its weights weigh at a chosen quantisation, and how much of a card the
// runtime keeps — and getting it wrong is not an error message, it is an
// out-of-memory twenty minutes into a load.
//
// It reads the model from its own repository and the hardware from the running
// fleet, so neither half is a number somebody retyped.
func newModelPlanCmd(client func() *ctlClient) *cobra.Command {
	var (
		context int
		params  int64
		precs   []string
		format  string
	)
	c := &cobra.Command{
		Use:   "plan <huggingface-repo>",
		Short: "Work out which of a model's files suits this cluster, and how it should be split",
		Long: "Lists every file the repository has — each GGUF quantisation, or with\n" +
			"--format safetensors its original weights — with the best way each could\n" +
			"run here: on one card, on several of a node's, over several nodes on a\n" +
			"fabric, or on a CPU node; what each card would hold and whether the cards\n" +
			"are free now. The model's shape is read from one file's header, nothing is\n" +
			"downloaded. `llmgw ctl model serve` without --quant takes the first.\n\n" +
			"With --precision it plans in the abstract instead, from the model's\n" +
			"config: every way it could run at bf16, fp8, q8_0 or q4_k_m, fewest GPUs\n" +
			"first, and a catalog entry for the best one.\n\n" +
			"Every candidate is listed rather than one answer, because the choice is\n" +
			"not the planner's to make: with cards to spare you want the\n" +
			"full-precision fit, and with one card free you want to know a quantised\n" +
			"one exists.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			repo := args[0]
			// The files the repository has, against the fleet as it is —
			// unless a precision or a size is given to plan in the abstract.
			if len(precs) == 0 && params == 0 {
				return planFiles(cmd, client, repo, format, context)
			}
			m, err := modelplan.FetchModel(cmd.Context(), repo)
			if err != nil {
				return err
			}
			if params > 0 {
				m.Params = params
			}
			if m.Params == 0 {
				return fmt.Errorf("%s does not publish a parameter count; pass --params", repo)
			}

			var chosen []modelplan.Precision
			for _, name := range precs {
				p, err := modelplan.PrecisionByName(name)
				if err != nil {
					return err
				}
				chosen = append(chosen, p)
			}

			var nodes []adminapi.NodeView
			if err := client().get(cmd.Context(), "/api/v1/nodes", &nodes); err != nil {
				return err
			}
			hw := hardwareFrom(nodes)
			if len(hw) == 0 {
				return fmt.Errorf("no node is ready, so there is nothing to plan against")
			}

			out := cmd.OutOrStdout()
			fmt.Fprintf(out, "%s: %.1fB parameters, %d layers, %d KV heads, trained to %d tokens\n",
				m.Name, float64(m.Params)/1e9, m.Layers, m.KVHeads, m.MaxContext)

			want := context
			if want <= 0 {
				want = m.MaxContext
			}
			fmt.Fprintf(out, "Planned for a %d-token pool. %s a token of context.\n\n",
				want, humanBytes(m.KVBytesPerToken(modelplan.BF16)))

			fits := modelplan.Fits(m, hw, want, chosen)
			if len(fits) == 0 {
				fmt.Fprintln(out, "It does not fit this cluster at any offered precision.")
				fmt.Fprintln(out, "\nWhat would change that:")
				fmt.Fprintln(out, "  - a smaller context pool, with --context")
				fmt.Fprintln(out, "  - a heavier quantisation, with --precision q4_k_m")
				fmt.Fprintln(out, "  - more cards free; llmgw ctl status shows what is in use")
				return errDoesNotFit
			}

			w := tab()
			fmt.Fprintln(w, "  GPUS\tPRECISION\tPER CARD\tFREE\tWHERE")
			for _, f := range fits {
				where := f.Class
				if f.CPUOnly {
					where = "cpu"
				} else if f.Nodes > 1 {
					where = fmt.Sprintf("%s, %d nodes on %s", f.Class, f.Nodes, f.Fabric)
				}
				fmt.Fprintf(w, "  %d\t%s\t%.1f GB\t%.1f GB\t%s\n",
					f.TP, f.Precision.Name, f.PerGPUGB, f.Headroom, where)
			}
			_ = w.Flush()

			fmt.Fprintf(out, "\nA catalog entry for the first of those:\n\n%s\n",
				catalogEntry(m, fits[0], repo, want))
			fmt.Fprintln(out, "Save it and add it with:  llmgw ctl model add -f <file>")
			return nil
		},
	}
	c.Flags().IntVar(&context, "context", 0, "token pool to plan for (default: the model's trained context)")
	c.Flags().Int64Var(&params, "params", 0, "parameter count, when the repository does not publish one")
	c.Flags().StringSliceVar(&precs, "precision", nil, "plan in the abstract, for these precisions (bf16, fp8, q8_0, q4_k_m), "+
		"from the model's config rather than its files")
	c.Flags().StringVar(&format, "format", "gguf", "gguf files, for llama.cpp, or safetensors, the original weights for vLLM")
	return c
}

// errDoesNotFit gives a script something to test without parsing prose.
var errDoesNotFit = fmt.Errorf("the model does not fit this cluster")

// hardwareFrom turns the fleet into pools a replica could actually occupy.
//
// Grouped by class and fabric together, because a replica may never span two
// fabrics: eight free cards split across two fabrics are not eight cards for a
// model that needs eight, and treating them as one pool would plan something
// that hangs on its first all-reduce.
func hardwareFrom(nodes []adminapi.NodeView) []modelplan.Hardware {
	type key struct{ class, fabric string }
	pools := map[key]*modelplan.Hardware{}

	for _, n := range nodes {
		if n.State != "ready" {
			continue
		}
		fabric := ""
		if len(n.Fabrics) > 0 {
			fabric = n.Fabrics[0]
		}
		// Each class of the node's devices is a pool of its own: a node with
		// a card in MIG mode has its instances' classes, and any whole card's
		// beside them. A node that does not count by class is its first
		// class, as it always was.
		kinds := n.ClassGPUs
		if len(kinds) == 0 {
			class := ""
			if len(n.Classes) > 0 {
				class = n.Classes[0]
			}
			kinds = []adminapi.ClassGPUs{{Class: class, Total: n.GPUTotal, Free: n.GPUFree}}
		}
		for _, kind := range kinds {
			class := kind.Class
			k := key{class, fabric}
			if pools[k] == nil {
				pools[k] = &modelplan.Hardware{Class: class, Fabric: fabric}
			}
			p := pools[k]
			p.Nodes++
			// The widest a single node could contribute, not the total: a
			// tensor group takes the same number of cards from every node it
			// spans.
			if kind.Free > p.GPUsPerNode {
				p.GPUsPerNode = kind.Free
			}
			if n.MemoryGB > p.HostMemGB {
				p.HostMemGB = n.MemoryGB
			}
			// Card memory from the class label, which carries it by
			// construction: "h100-80g" is eighty gigabytes and says so, and
			// "a100-80g-mig-3g.40gb" is an instance of forty.
			if class != "" && p.GPUMemGB == 0 {
				p.GPUMemGB = memFromClass(class)
			}
		}
	}

	out := make([]modelplan.Hardware, 0, len(pools))
	for _, p := range pools {
		if p.Class != "" && (p.GPUMemGB == 0 || p.GPUsPerNode == 0) {
			continue // a GPU pool with no free cards is not a pool
		}
		out = append(out, *p)
	}
	return out
}

// memFromClass reads the card size out of a class label such as "h100-80g".
func memFromClass(class string) int { return modelplan.MemFromClass(class) }

// catalogEntry renders a fit as the YAML that would serve it.
func catalogEntry(m modelplan.Model, f modelplan.Fit, source string, context int) string {
	var b strings.Builder
	name := strings.ToLower(m.Name)
	fmt.Fprintf(&b, "name: %s\n", name)
	if f.Precision.GGUF {
		// A GGUF quantisation lives in a file, not a repository, and llama.cpp
		// is given a path. Printing the repository id here would produce an
		// entry that looks right and cannot load — which is the failure this
		// whole command exists to prevent.
		fmt.Fprintf(&b, "# %s is a GGUF quantisation: download the %s file and give its path.\n",
			f.Precision.Name, f.Precision.Name)
		fmt.Fprintf(&b, "source: /var/lib/llmgw-models/%s-%s.gguf\n", name, f.Precision.Name)
	} else {
		fmt.Fprintf(&b, "source: %s\n", source)
	}
	fmt.Fprintf(&b, "routing_profile: chat\n")
	fmt.Fprintf(&b, "context_length: %d\n", context)
	fmt.Fprintf(&b, "scaling: { min: 1, max: 1 }\n")
	fmt.Fprintf(&b, "profiles:\n")

	switch {
	case f.CPUOnly:
		fmt.Fprintf(&b, "  - name: cpu\n    engine: %s\n", f.Precision.Engine)
		fmt.Fprintf(&b, "    min_memory_gb: %d\n", int(f.PerGPUGB)+2)
	default:
		fmt.Fprintf(&b, "  - name: tp%d\n    engine: %s\n", f.TP, f.Precision.Engine)
		fmt.Fprintf(&b, "    gpu_class: %s\n    tp: %d\n", f.Class, f.TP)
		if f.Nodes > 1 {
			fmt.Fprintf(&b, "    nodes: %d\n    gpus_per_node: %d\n", f.Nodes, f.GPUsPerNode)
		}
		if q := f.Precision.Quantization; q != "" {
			fmt.Fprintf(&b, "    args: [\"--quantization\", \"%s\"]\n", q)
		}
	}
	fmt.Fprintf(&b, "    priority: 100\n    weight: 1\n")
	// Measured rather than guessed is the rule everywhere else in this system,
	// and a planner should say so rather than pretend its estimate is one.
	fmt.Fprintf(&b, "    startup_timeout: 10m\n")
	b.WriteString("\n# weight should come from `llmgw ctl bench` once it is running;\n")
	b.WriteString("# 1 is a placeholder, not a measurement.\n")
	return b.String()
}

// humanBytes renders a per-token cache size, which is always small.
func humanBytes(n int64) string {
	switch {
	case n >= 1<<20:
		return fmt.Sprintf("%.1f MiB", float64(n)/(1<<20))
	default:
		return fmt.Sprintf("%.0f KiB", float64(n)/1024)
	}
}

// planFiles prints the control plane's plan of a repository's files.
func planFiles(cmd *cobra.Command, client func() *ctlClient, repo, format string, context int) error {
	q := url.Values{"repo": {repo}, "format": {format}}
	if context > 0 {
		q.Set("context", strconv.Itoa(context))
	}
	var p modelserve.Plan
	if err := client().get(cmd.Context(), "/api/v1/plan?"+q.Encode(), &p); err != nil {
		return err
	}
	out := cmd.OutOrStdout()
	fmt.Fprintf(out, "%s at %s", p.Repo, catalog.ShortRevision(p.Revision))
	if p.Arch != "" {
		fmt.Fprintf(out, ", %s", p.Arch)
	}
	if p.Params > 0 {
		fmt.Fprintf(out, ", %.1fB parameters", float64(p.Params)/1e9)
	}
	fmt.Fprintf(out, "; planned for a %d-token context\n", p.Context)
	if p.Gated {
		fmt.Fprintln(out, "gated: the mirror cannot fetch it without an account's approval")
	}
	if !p.Servable {
		fmt.Fprintf(out, "the nodes' engine cannot run it: %s\n", p.Reason)
		return errDoesNotFit
	}
	if p.Note != "" {
		fmt.Fprintln(out, "note: "+p.Note)
	}
	fmt.Fprintln(out)
	w := tabTo(out)
	fmt.Fprintln(w, "FILE\tSIZE\tBEST WAY\tHOLDS\tOF\tNOW")
	for _, a := range p.Artifacts {
		name := a.Name
		if a.Shards > 1 && a.Format == "gguf" {
			name += fmt.Sprintf(" (%d parts)", a.Shards)
		}
		if a.Best == nil {
			fmt.Fprintf(w, "%s\t%s\tfits nowhere\t\t\t\n", name, bytesShort(a.Bytes))
			continue
		}
		b := a.Best
		now := "busy"
		if b.FreeNow {
			now = "free"
		}
		holds := fmt.Sprintf("%.1f GB a card", b.NeedGB)
		if b.CPU {
			holds = fmt.Sprintf("%.1f GB of RAM", b.NeedGB)
		}
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%.0f GB\t%s\n", name, bytesShort(a.Bytes), b.Label(), holds, b.HaveGB, now)
	}
	_ = w.Flush()
	rec, ok := p.Recommended()
	if !ok {
		fmt.Fprintln(out, "\nNo file fits this cluster as it is: a smaller context (--context), or more memory.")
		return errDoesNotFit
	}
	fmt.Fprintf(out, "\nBest: %s on %s.\n", rec.Artifact, rec.Label())
	if len(p.Options) > 1 {
		fmt.Fprintln(out, "\nEvery way, best first:")
		w = tabTo(out)
		for _, o := range p.Options {
			now := "busy"
			if o.FreeNow {
				now = "free"
			}
			fmt.Fprintf(w, "  %s\t%s\t%.1f of %.0f GB\t%d host(s)\t%s\n", o.Artifact, o.Label(), o.NeedGB, o.HaveGB, o.Hosts, now)
		}
		_ = w.Flush()
	}
	fmt.Fprintf(out, "\nServe it: llmgw ctl model serve %s", p.Repo)
	if p.Format == "safetensors" {
		fmt.Fprint(out, " --format safetensors")
	}
	fmt.Fprintln(out, " --fallbacks")
	return nil
}