package cli

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"runtime"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/release"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
)

func newUpgradeCmd(client func() *ctlClient) *cobra.Command {
	var (
		nodes     []string
		skipCP    bool
		releases  string
		perNode   time.Duration
		serverSvc string
		parallel  int
		sshUser   string
		othersCP  bool
	)
	c := &cobra.Command{
		Use:   "upgrade VERSION",
		Short: "Upgrade the control planes, then every node, in waves",
		Long: "Moves the cluster to a release: v0.10.4, or a candidate like v0.10.4-rc.2.\n\n" +
			"The other control planes first, over ssh as --ssh-user (the one sudo was\n" +
			"run by), then this machine — run it with sudo on a control plane, since it\n" +
			"replaces /usr/local/bin/llmgw and restarts llmgw-server. A leader restarted\n" +
			"last hands over to a standby already on the release.\n\n" +
			"Then the nodes, in waves: one first, then twice as many each time a wave\n" +
			"comes back whole, up to --parallel at once. Each agent fetches the release\n" +
			"for its own platform, checks it against the release's checksums, swaps it\n" +
			"in and restarts; engines keep serving through all of it and are adopted\n" +
			"again, and no ssh into the nodes is needed.\n\n" +
			"A node counts as done when it is back on the new version, ready, with as\n" +
			"many ready replicas as before. A wave not back within --timeout stops the\n" +
			"upgrade there, so a bad release reaches one node rather than all of them.\n" +
			"Nodes already on the version are skipped, so running it again resumes.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			v := args[0]
			if !release.Tag.MatchString(v) {
				return fmt.Errorf("%q is not a release version (v1.2.3 or v1.2.3-rc.4)", v)
			}
			out := cmd.OutOrStdout()
			ctx := cmd.Context()
			if releases == "" {
				// Where the control plane sends its nodes: an offline
				// cluster's mirror, so this machine needs no internet either.
				if cfg, ok := loadServerConfig("/etc/llmgw/server.yaml"); ok {
					releases = cfg.ReleasesBase()
				}
				if releases == "" {
					releases = release.DefaultBase
				}
			}
			if !skipCP {
				// The leader last, wherever this runs: every other control
				// plane is a standby, upgraded with nothing moving, and the
				// leader's restart hands over once, to a machine already on
				// the release.
				others := func() error {
					if !othersCP {
						return nil
					}
					return upgradeOtherControlPlanes(ctx, out, client, v, releases, sshUser)
				}
				self := func() error { return upgradeControlPlane(ctx, out, client, v, releases, serverSvc) }
				order := []func() error{others, self}
				if !thisLeads(ctx, client) {
					order = []func() error{self, others}
				}
				for _, step := range order {
					if err := step(); err != nil {
						return err
					}
				}
			}
			return upgradeNodes(ctx, out, client, v, nodes, perNode, parallel)
		},
	}
	c.Flags().StringSliceVar(&nodes, "nodes", nil, "only these nodes (default: all)")
	c.Flags().BoolVar(&skipCP, "skip-control-plane", false, "leave this machine as it is and upgrade only nodes")
	c.Flags().StringVar(&releases, "releases", "", "where release binaries are fetched from, for this machine "+
		"(default: as server.yaml says, the mirror's copy on an offline cluster, else the project's)")
	c.Flags().DurationVar(&perNode, "timeout", 3*time.Minute, "how long a node may take to come back")
	c.Flags().StringVar(&serverSvc, "service", "llmgw-server", "this machine's systemd unit")
	c.Flags().IntVar(&parallel, "parallel", 20, "the most nodes upgraded at once, once the first waves are back")
	c.Flags().StringVar(&sshUser, "ssh-user", os.Getenv("SUDO_USER"), "who to log in to the other control planes as")
	c.Flags().BoolVar(&othersCP, "other-control-planes", true, "upgrade the cluster's other control planes too, over ssh")
	return c
}

// thisLeads says whether this machine's control plane leads; one that does
// not say, alone or older, is taken to.
func thisLeads(ctx context.Context, client func() *ctlClient) bool {
	var sum adminapi.Summary
	if err := client().get(ctx, "/api/v1/summary", &sum); err != nil {
		return true
	}
	for _, cp := range sum.ControlPlanes.Members {
		if cp.This {
			return cp.State == "leading"
		}
	}
	return true
}

// upgradeOtherControlPlanes moves every other control plane of the cluster to
// v, over ssh with `llmgw upgrade`. Nothing is changed here if one cannot be
// reached.
func upgradeOtherControlPlanes(ctx context.Context, out io.Writer, client func() *ctlClient, v, releases, user string) error {
	var sum adminapi.Summary
	if err := client().get(ctx, "/api/v1/summary", &sum); err != nil {
		return err
	}
	for _, cp := range sum.ControlPlanes.Members {
		if cp.This {
			continue
		}
		r := &onboard.SSH{Addr: cp.Address, User: user}
		have, err := r.Run(ctx, "llmgw --version", nil)
		if err != nil {
			return fmt.Errorf("%s (%s) cannot be reached over ssh as %q: %v\n"+
				"run `sudo llmgw upgrade %s --service llmgw-server` there, then this again with "+
				"--other-control-planes=false", cp.Name, cp.Address, user, err, v)
		}
		if strings.Contains(string(have), " "+v+"\n") || strings.HasSuffix(strings.TrimSpace(string(have)), " "+v) {
			fmt.Fprintf(out, "== %s already runs %s\n", cp.Name, v)
			continue
		}
		fmt.Fprintf(out, "== %s (%s): %s -> %s\n", cp.Name, cp.State, strings.TrimSpace(string(have)), v)
		res, err := r.Run(ctx, "sudo llmgw upgrade "+shellQuote(v)+" --service llmgw-server --releases "+shellQuote(releases), nil)
		if err != nil {
			return fmt.Errorf("upgrading %s: %v: %s", cp.Name, err, strings.TrimSpace(string(res)))
		}
		fmt.Fprintf(out, "   upgraded and restarted\n")
	}
	return nil
}

// controlPlaneStep is what this machine needs to run v: the binary installed,
// the service restarted, both, or neither. The binary is this command's own;
// running is what the server says it runs, "" from one too old to say. A
// binary put in place by hand is not a server running it.
func controlPlaneStep(binary, running, v string) (install, restart bool) {
	install = binary != v
	return install, install || running != v
}

func upgradeControlPlane(ctx context.Context, out io.Writer, client func() *ctlClient, v, releases, svc string) error {
	var sum adminapi.Summary
	_ = client().get(ctx, "/api/v1/summary", &sum)
	install, restart := controlPlaneStep(version.Version, sum.Version, v)
	if !restart {
		fmt.Fprintf(out, "this machine already runs %s\n", v)
		return nil
	}
	if os.Geteuid() != 0 {
		return errors.New("replacing this machine's binary needs root: run it with sudo, " +
			"or pass --skip-control-plane to upgrade only the nodes")
	}
	self, err := os.Executable()
	if err != nil {
		return err
	}
	backup := "the binary installed by hand before it"
	if install {
		fmt.Fprintf(out, "== this machine: %s -> %s\n", version.Version, v)
		if backup, err = release.Install(ctx, nil, releases, v, runtime.GOOS, runtime.GOARCH, self); err != nil {
			return err
		}
		fmt.Fprintf(out, "   installed, checksum verified; the old binary is %s\n", backup)
	} else {
		running := sum.Version
		if running == "" {
			running = "an older release"
		}
		fmt.Fprintf(out, "== this machine: %s is installed, and %s still runs it; restarting %s\n", v, running, svc)
	}
	var before []adminapi.NodeView
	_ = client().get(ctx, "/api/v1/nodes", &before)
	if b, err := exec.CommandContext(ctx, "systemctl", "restart", svc).CombinedOutput(); err != nil { //nolint:gosec // fixed command
		return fmt.Errorf("restarting %s: %v: %s", svc, err, strings.TrimSpace(string(b)))
	}
	fmt.Fprintf(out, "   %s restarted; waiting for its nodes to reconnect\n", svc)
	// Back when the admin API answers and every node that was ready is ready
	// again: agents reconnect within seconds of the control plane returning.
	want := 0
	for _, n := range before {
		if n.State == "ready" {
			want++
		}
	}
	deadline := time.Now().Add(2 * time.Minute)
	for {
		var now []adminapi.NodeView
		if err := client().get(ctx, "/api/v1/nodes", &now); err == nil {
			ready := 0
			for _, n := range now {
				if n.State == "ready" {
					ready++
				}
			}
			if ready >= want {
				fmt.Fprintf(out, "   back: %d of %d nodes ready\n", ready, len(now))
				return nil
			}
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("%s did not come back with its %d nodes within two minutes; "+
				"see journalctl -u %s, and %s holds the previous binary", svc, want, svc, backup)
		}
		time.Sleep(2 * time.Second)
	}
}

// upgradeNodes moves nodes to v in waves: one node first, then twice as
// many each time one comes back whole, up to parallel at once.
//
// Waves rather than one at a time, because one at a time is hours on a
// fleet of hundreds, and an agent's restart costs nothing served — the
// engines keep running and are adopted again. One first, because a release
// that does not come back should reach one node rather than a wave of them.
func upgradeNodes(ctx context.Context, out io.Writer, client func() *ctlClient, v string, only []string,
	perNode time.Duration, parallel int,
) error {
	var nodes []adminapi.NodeView
	if err := client().get(ctx, "/api/v1/nodes", &nodes); err != nil {
		return err
	}
	slices.SortFunc(nodes, func(a, b adminapi.NodeView) int { return strings.Compare(a.ID, b.ID) })
	var todo []adminapi.NodeView
	var away []string
	skipped := 0
	for _, n := range nodes {
		if len(only) > 0 && !slices.Contains(only, n.ID) {
			continue
		}
		if n.AgentVersion == v {
			skipped++
			continue
		}
		// A node that is not connected cannot be told, and one tried first
		// stopped every run at it: passed over, and named at the end.
		if n.State == "lost" || n.State == "unreachable" {
			away = append(away, n.ID)
			continue
		}
		todo = append(todo, n)
	}
	if len(away) > 0 {
		fmt.Fprintf(out, "Passed over, not connected: %s\n", strings.Join(away, ", "))
	}
	parallel = max(parallel, 1)
	start, size, done := time.Now(), 1, 0
	for len(todo) > 0 {
		wave := todo[:min(size, len(todo))]
		todo = todo[len(wave):]
		if err := upgradeWave(ctx, out, client, v, wave, perNode); err != nil {
			return fmt.Errorf("%w\nstopped after this wave: %d upgraded, %d not tried", err, done, len(todo))
		}
		done += len(wave)
		size = min(size*2, parallel)
	}
	fmt.Fprintf(out, "%d node(s) upgraded to %s in %s, %d already on it\n",
		done, v, time.Since(start).Round(time.Second), skipped)
	// Checked at the end, against what the control plane says now: a node
	// that is not on the version is said by name, and the command fails. A
	// site's nodes stayed a release behind with nothing printed to say so.
	if len(nodes) == 0 {
		fmt.Fprintln(out, "This control plane has no nodes registered, so none were upgraded.")
	}
	var now []adminapi.NodeView
	if err := client().get(ctx, "/api/v1/nodes", &now); err == nil {
		var behind []string
		for _, n := range now {
			if (len(only) == 0 || slices.Contains(only, n.ID)) && n.AgentVersion != v {
				behind = append(behind, fmt.Sprintf("%s (%s, %s)", n.ID, orUnknown(n.AgentVersion), n.State))
			}
		}
		if len(behind) > 0 {
			slices.Sort(behind)
			return fmt.Errorf("not on %s: %s\nrun this again to try them, or `llmgw ctl logs NODE` for why",
				v, strings.Join(behind, ", "))
		}
	}
	// The one machine this does not reach. Said every time, because it is
	// exactly the kind of thing that stays three versions behind unnoticed.
	fmt.Fprintf(out, "The model mirror's host is not upgraded by this; on it, run:\n"+
		"  sudo llmgw upgrade %s --service llmgw-mirror\n", v)
	return nil
}

// upgradeWave tells every node in the wave at once, and waits until each is
// back on v and serving what it served.
//
// Told together, not one after another: each answer can take minutes, and a
// wave of twenty waited for them in turn with nothing on the screen.
func upgradeWave(ctx context.Context, out io.Writer, client func() *ctlClient, v string,
	wave []adminapi.NodeView, perNode time.Duration,
) error {
	start := time.Now()
	before := readyReplicasByNode(ctx, client)
	ids := make([]string, len(wave))
	for i, n := range wave {
		ids[i] = n.ID
	}
	fmt.Fprintf(out, "== %d node(s): %s\n", len(wave), abbreviate(ids, 8))
	body, _ := json.Marshal(adminapi.UpgradeRequest{Version: v})
	var (
		mu     sync.Mutex
		failed []string
		wg     sync.WaitGroup
	)
	for _, n := range wave {
		wg.Add(1)
		go func() {
			defer wg.Done()
			// The control plane waits for the agent's answer, or for it to
			// come back on v, longer than the agent gives its download.
			err := client().patient(agentsvc.UpgradeWait+time.Minute).post(ctx, "/api/v1/nodes/"+n.ID+"/upgrade", body, nil)
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				failed = append(failed, fmt.Sprintf("%s: %v", n.ID, err))
				fmt.Fprintf(out, "   %s: not upgraded: %v\n", n.ID, err)
				return
			}
			fmt.Fprintf(out, "   %s: on %s, coming back\n", n.ID, v)
		}()
	}
	wg.Wait()
	waiting := slices.DeleteFunc(slices.Clone(ids), func(id string) bool {
		return slices.ContainsFunc(failed, func(f string) bool { return strings.HasPrefix(f, id+": ") })
	})
	deadline := time.Now().Add(perNode)
	backAt := map[string]time.Time{}
	said := time.Now()
	for len(waiting) > 0 {
		var now []adminapi.NodeView
		_ = client().get(ctx, "/api/v1/nodes", &now)
		var reps []adminapi.ReplicaView
		_ = client().get(ctx, "/api/v1/replicas", &reps)
		after, starting := replicasByNode(reps)
		state := map[string]adminapi.NodeView{}
		for _, x := range now {
			state[x.ID] = x
		}
		waiting = slices.DeleteFunc(waiting, func(id string) bool {
			cur := state[id]
			// Connected and on v: ready, or cordoned or draining as it
			// was, which an upgrade does not undo.
			if cur.AgentVersion != v || (cur.State != "ready" && cur.State != "cordoned" && cur.State != "draining") {
				return false
			}
			if _, ok := backAt[id]; !ok {
				backAt[id] = time.Now()
			}
			// Serving what it served; or, a while on, nothing on it still
			// starting: a model that went idle meanwhile is not one lost.
			return after[id] >= before[id] || (time.Since(backAt[id]) > 90*time.Second && starting[id] == 0)
		})
		if len(waiting) == 0 {
			break
		}
		if time.Now().After(deadline) {
			var why []string
			for _, id := range waiting {
				cur := state[id]
				why = append(why, fmt.Sprintf("%s (version %q, state %q, replicas %d/%d)",
					id, cur.AgentVersion, cur.State, after[id], before[id]))
			}
			failed = append(failed, fmt.Sprintf("not back within %s: %s", perNode, strings.Join(why, ", ")))
			break
		}
		if time.Since(said) >= 30*time.Second {
			var w []string
			for _, id := range waiting {
				cur := state[id]
				w = append(w, fmt.Sprintf("%s (%s, %s, replicas %d/%d)", id, orUnknown(cur.AgentVersion), cur.State,
					after[id], before[id]))
			}
			fmt.Fprintf(out, "   waiting for %s\n", abbreviate(w, 8))
			said = time.Now()
		}
		time.Sleep(2 * time.Second)
	}
	if len(failed) > 0 {
		return fmt.Errorf("not upgraded:\n  %s\n`llmgw ctl logs NODE` says why; each keeps its previous binary beside the new one",
			strings.Join(failed, "\n  "))
	}
	fmt.Fprintf(out, "   done in %s\n", time.Since(start).Round(time.Second))
	return nil
}

// replicasByNode counts each node's ready replicas, and those still placing
// or starting.
func replicasByNode(reps []adminapi.ReplicaView) (ready, starting map[string]int) {
	ready, starting = map[string]int{}, map[string]int{}
	for _, r := range reps {
		seen := map[string]bool{}
		for _, sh := range r.Shards {
			if seen[sh.Node] {
				continue
			}
			seen[sh.Node] = true
			switch r.State {
			case "ready":
				ready[sh.Node]++
			case "placing", "starting":
				starting[sh.Node]++
			}
		}
	}
	return ready, starting
}

// abbreviate lists up to n names, and says how many more.
func abbreviate(names []string, n int) string {
	if len(names) <= n {
		return strings.Join(names, ", ")
	}
	return strings.Join(names[:n], ", ") + fmt.Sprintf(" and %d more", len(names)-n)
}

// readyReplicasByNode counts, for each node, the ready replicas with a
// shard on it.
func readyReplicasByNode(ctx context.Context, client func() *ctlClient) map[string]int {
	var reps []adminapi.ReplicaView
	out := map[string]int{}
	if err := client().get(ctx, "/api/v1/replicas", &reps); err != nil {
		return out
	}
	for _, r := range reps {
		if r.State != "ready" {
			continue
		}
		seen := map[string]bool{}
		for _, sh := range r.Shards {
			if !seen[sh.Node] {
				seen[sh.Node] = true
				out[sh.Node]++
			}
		}
	}
	return out
}