package web

import (
	"context"
	"errors"
	"fmt"
	"gitea.mynodes.xyz/ilker/llmgw/internal/gguf"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
)

// CatalogAdmin adds a model to the catalogue in force.
type CatalogAdmin interface {
	Add(ctx context.Context, m catalog.Model, by string) error
	Declared(name string) bool
	// Sources maps a weights source to the catalog models that use it.
	Sources(ctx context.Context) (map[string][]catalog.Model, error)
	// Scale changes how a model added at runtime runs: always, on demand, or
	// not at all. One declared in configuration is refused.
	Scale(ctx context.Context, name string, sc catalog.Scaling, by string) (catalog.Model, error)
}

// serveForm is the offer to put one mirrored version into service.
//
// Everything a catalogue entry needs that the mirror already knows is filled
// in: the name, the source, and the engine that reads that kind of file. What
// is left is what nobody can infer — the context pool and how many replicas —
// and those are fields rather than guesses presented as facts.
type serveForm struct {
	Model   string
	Version string
	Source  string
	Engine  string
	// Profile is the profile name, which doubles as how a node_selector would
	// be written later. Named for the hardware rather than the model, because
	// a profile is about where a model runs.
	Profile string
	// Files is how many files the version holds, above one for a set.
	Files int
	// MinMemoryGB is the host RAM to ask for, estimated from the weights.
	//
	// Filled in rather than left blank. Blank writes no constraint at all, and
	// a CPU profile has no card to be short of, so every ready node becomes a
	// candidate: a 4.7 GB model went onto a 7.9 GB machine and left it with
	// 3.5 GB for a KV pool. Nothing was wrong with placement — it was never
	// told what the model needed.
	MinMemoryGB int

	// Refusal is set when this version cannot be served as it stands, and says
	// why. Offered and then failing at placement is worse than not offered.
	Refusal string
	// Warning is set when it can be served but probably will not run.
	Warning string
	// Already is true when a model of this name is in the catalogue.
	Already bool
}

// planServe works out what serving this version would mean.
func planServe(st mirror.Status, files []mirror.Status, base string,
	fleet Source, cat CatalogAdmin, check func(source string) gguf.Verdict,
) serveForm {
	f := serveForm{
		Model:   st.Model,
		Version: st.Version,
		Source: mirror.Entry{Model: st.Model, Version: st.Version, File: st.File}.
			URL(base),
		Engine:  engineFor(st.File),
		Profile: "cpu",
	}

	// A version of several files is served as the version, not as one of them.
	// The source is the directory — a trailing slash — and the agent fetches
	// every file in it and hands the engine the folder.
	if len(files) > 1 {
		f.Source = strings.TrimSuffix(
			mirror.Entry{Model: st.Model, Version: st.Version}.URL(base), "/") + "/"
		f.Files = len(files)
		// From whichever file says, because a version is a mixture only when
		// somebody has put one there by hand. The weights decide the engine.
		for _, one := range files {
			if e := engineFor(one.File); e != "" {
				f.Engine = e
				break
			}
		}
		var total int64
		for _, one := range files {
			total += one.BytesDone
		}
		st.BytesTotal = total
	}

	switch f.Engine {
	case "":
		f.Refusal = "nothing here reads " + extOf(st.File) + " files"
		return f
	case "llamacpp":
		f.Profile = "cpu"
		if check != nil {
			if v := check(f.Source); !v.Servable {
				f.Refusal = "llama.cpp on the nodes cannot run this: " + v.Reason
				return f
			}
		}
	default:
		f.Profile = "gpu"
	}
	// Whether any node can actually run it. Agents report their installed
	// engines now, so this is a fact rather than the guess it started as —
	// which was cards, and would have passed a GPU node with no vLLM on it.
	//
	// A warning, not a refusal: an engine can be installed this afternoon, and
	// a node that reports nothing is an older agent rather than an empty one.
	if !runsAnywhere(fleet, f.Engine) {
		f.Warning = "no node reports having " + f.Engine + " installed, and " +
			extOf(st.File) + " weights need it"
		if f.Engine != "llamacpp" && gpus(fleet) == 0 {
			f.Warning += ". This fleet has no GPU either; a GGUF build of the " +
				"same model runs on the CPU nodes"
		}
	}
	f.MinMemoryGB = memoryFor(st.BytesTotal, st.BytesDone)
	if biggest := biggestNode(fleet); biggest > 0 && f.MinMemoryGB > biggest {
		f.Warning = strings.TrimSuffix(f.Warning, ".") +
			fmt.Sprintf("%sthe largest node has %d GB and this needs about %d",
				sep(f.Warning), biggest, f.MinMemoryGB)
	}
	if cat != nil && cat.Declared(st.Model) {
		f.Already = true
		f.Refusal = "a model called " + st.Model +
			" is declared in configuration; edit it there rather than adding a second"
	}
	return f
}

// engineFor is the engine that reads this kind of weights file.
func engineFor(file string) string {
	switch strings.ToLower(extOf(file)) {
	case ".gguf":
		return "llamacpp"
	case ".safetensors", ".bin":
		return "vllm"
	}
	return ""
}

func extOf(file string) string {
	if i := strings.LastIndex(file, "."); i >= 0 {
		return file[i:]
	}
	return file
}

// runsAnywhere reports whether some node has this engine.
//
// True when nothing is known, so a fleet of agents older than the engine probe
// is not told every model is unservable. That matches what placement does with
// the same absence, and the two disagreeing would be worse than either.
func runsAnywhere(fleet Source, engine string) bool {
	if fleet == nil || engine == "" {
		return true
	}
	anyReported := false
	for _, n := range fleet.Nodes() {
		if len(n.Engines) == 0 {
			continue
		}
		anyReported = true
		for _, e := range n.Engines {
			if e == engine {
				return true
			}
		}
	}
	return !anyReported
}

func gpus(fleet Source) int {
	if fleet == nil {
		return 0
	}
	n := 0
	for _, node := range fleet.Nodes() {
		n += node.GPUTotal
	}
	return n
}

// buildModel turns the submitted form into a catalogue entry.
func buildModel(get func(string) string) (catalog.Model, error) {
	m := catalog.Model{
		Name:           strings.TrimSpace(get("name")),
		Description:    strings.TrimSpace(get("description")),
		Source:         strings.TrimSpace(get("source")),
		RoutingProfile: "chat",
	}
	if m.Name == "" || m.Source == "" {
		return m, errors.New("a model needs a name and a source")
	}
	for _, a := range strings.Split(get("aliases"), ",") {
		if a = strings.TrimSpace(a); a != "" {
			m.Aliases = append(m.Aliases, a)
		}
	}

	ctxLen, err := positive(get("context_length"), "the context length")
	if err != nil {
		return m, err
	}
	m.ContextLength = ctxLen
	if v := strings.TrimSpace(get("max_output")); v != "" {
		if m.MaxOutput, err = positive(v, "the output limit"); err != nil {
			return m, err
		}
	}

	min, err := positive(get("replicas"), "the replica count")
	if err != nil {
		return m, err
	}
	m.Scaling = scalingFor(get("mode"), min, 0)

	engine := strings.TrimSpace(get("engine"))
	if engine == "" {
		return m, errors.New("which engine should run it?")
	}
	p := catalog.Profile{
		Name: strings.TrimSpace(get("profile")), Engine: engine,
		Weight: 1, Priority: 100,
	}
	if p.Name == "" {
		p.Name = engine
	}
	// The pool, not the per-request limit: llama.cpp divides it across its
	// parallel slots, so the number a client may use is smaller and is what
	// context_length above says.
	if v := strings.TrimSpace(get("ctx_size")); v != "" {
		if _, err := positive(v, "the engine's context pool"); err != nil {
			return m, err
		}
		p.Args = []string{"--ctx-size", v}
	}
	if v := strings.TrimSpace(get("min_memory_gb")); v != "" {
		if p.MinMemoryGB, err = positive(v, "the memory requirement"); err != nil {
			return m, err
		}
	}
	if k, v := strings.TrimSpace(get("selector_key")), strings.TrimSpace(get("selector_value")); k != "" && v != "" {
		// What keeps a large model off the small nodes. A CPU profile has no
		// card to be short of, so without this every ready node looks like a
		// candidate and the first one picked is the one that gets OOM-killed.
		p.NodeSelector = map[string]string{k: v}
	}
	m.Profiles = []catalog.Profile{p}
	return m, nil
}

func positive(v, what string) (int, error) {
	n, err := strconv.Atoi(strings.TrimSpace(v))
	if err != nil || n < 0 {
		return 0, fmt.Errorf("%s must be a whole number", what)
	}
	return n, nil
}

// memoryFor estimates the host RAM a model needs, from the size of its weights.
//
// Weights plus half again, plus two. The half covers the KV pool, which grows
// with the same things the weights do — layers and heads — and the two are the
// engine's own footprint and the page cache it will not get back. It is an
// estimate and the field it fills is editable; what it is not is blank, which
// is what put a 4.7 GB model on a 7.9 GB node.
//
// The honest number comes from running it: `llmgw ctl bench` measures what a
// replica can actually do, and that is a different question from what it needs
// to start.
func memoryFor(total, done int64) int {
	size := total
	if size <= 0 {
		size = done
	}
	if size <= 0 {
		return 0
	}
	gb := float64(size) / (1 << 30)
	return int(gb*1.5) + 2
}

// biggestNode is the most memory any single node reports, or 0 when none says.
func biggestNode(fleet Source) int {
	if fleet == nil {
		return 0
	}
	most := 0
	for _, n := range fleet.Nodes() {
		if n.MemoryGB > most {
			most = n.MemoryGB
		}
	}
	return most
}

// sep joins a second clause onto a warning that may be empty.
func sep(existing string) string {
	if existing == "" {
		return ""
	}
	return ", and "
}

// Run modes a page offers for a model, in the words a person would use.
const (
	modeAlways   = "always"
	modeOnDemand = "ondemand"
	modeStopped  = "stopped"
)

// defaultIdle is how long an on-demand model set from a page stays loaded
// after its last request.
const defaultIdle = 15 * time.Minute

// scalingFor turns a run mode into scaling: replicas is how many run, or may
// run on demand; idle keeps an on-demand model's timeout when it has one.
func scalingFor(mode string, replicas int, idle time.Duration) catalog.Scaling {
	if replicas < 1 {
		replicas = 1
	}
	switch mode {
	case modeOnDemand:
		if idle <= 0 {
			idle = defaultIdle
		}
		return catalog.Scaling{Min: 0, Max: replicas, IdleTimeout: idle}
	case modeStopped:
		return catalog.Scaling{Min: 0, Max: 0}
	default:
		return catalog.Scaling{Min: replicas, Max: replicas}
	}
}

// modeOf names how a model runs.
func modeOf(sc catalog.Scaling) string {
	switch {
	case sc.Min > 0:
		return modeAlways
	case sc.Max > 0:
		return modeOnDemand
	default:
		return modeStopped
	}
}