package cli

import (
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/spf13/cobra"

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

func newModelServeCmd(client func() *ctlClient) *cobra.Command {
	var (
		req      modelserve.Request
		onDemand bool
		detach   bool
	)
	c := &cobra.Command{
		Use:   "serve OWNER/REPO",
		Short: "Take a model from the registry to a replica that answers, in one command",
		Long: "Picks the GGUF file (--quant, Q4_K_M by default, or --file), has the mirror\n" +
			"fetch it and checks what arrived against the registry's checksum, checks\n" +
			"that the nodes' llama.cpp can load it and sizes its memory from its own\n" +
			"header, adds it to the catalog, waits for a replica, and asks it to say\n" +
			"hello. Each step is shown as it happens; the first that fails stops it.\n\n" +
			"Without --quant or --file, the file is the one that suits the cluster best,\n" +
			"and where it runs is the plan's best way for it: on one card, on several of\n" +
			"a node's, over several nodes on a fabric, or on a CPU node (`llmgw ctl model\n" +
			"plan REPO` lists them). --cpu or --gpu-class with --gpus and --nodes say it\n" +
			"instead; --fallbacks adds the next best ways for when the first is busy.\n\n" +
			"A vision model's projector (--projector) is fetched beside it and loaded with\n" +
			"it. --format safetensors serves the original weights instead, with their\n" +
			"config.json and tokenizer, on vLLM.\n\n" +
			"  llmgw ctl model serve bartowski/Qwen2.5-3B-Instruct-GGUF\n" +
			"  llmgw ctl model serve Qwen/Qwen2.5-7B-Instruct-GGUF --quant Q5_K_M --name qwen-7b\n" +
			"  llmgw ctl model serve owner/big-GGUF --on-demand\n" +
			"  llmgw ctl model serve google/gemma-3-4b-it-GGUF --projector mmproj-model-f16.gguf\n" +
			"  llmgw ctl model serve Qwen/Qwen2.5-7B-Instruct --format safetensors --fallbacks\n" +
			"  llmgw ctl model serve meta-llama/Llama-3.1-405B-Instruct --format safetensors --gpu-class h100-80g --gpus 8 --nodes 2",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			req.Repo = args[0]
			if onDemand {
				req.Replicas = 0
			}
			body, _ := json.Marshal(req)
			var job modelserve.Job
			if err := client().post(cmd.Context(), "/api/v1/models/serve", body, &job); err != nil {
				return err
			}
			if detach {
				return detached(cmd, job)
			}
			return followJob(cmd, client, job, "is served")
		},
	}
	f := c.Flags()
	f.StringVar(&req.Quant, "quant", "", "the quantisation to serve (default: the one that suits the cluster best)")
	f.StringVar(&req.File, "file", "", "the file to serve, instead of choosing by --quant")
	f.StringVar(&req.Projector, "projector", "", "a vision model's mmproj file, fetched and loaded beside it")
	f.StringVar(&req.Format, "format", "gguf", "gguf, run by llama.cpp, or safetensors: the original weights, run by vLLM on GPUs")
	f.BoolVar(&req.CPU, "cpu", false, "run it on the CPU, whatever the plan says")
	f.StringVar(&req.GPUClass, "gpu-class", "", "run it on GPUs of this class, whatever the plan says")
	f.IntVar(&req.GPUsPerNode, "gpus", 0, "with --gpu-class: the cards one replica takes on each node (default 1)")
	f.IntVar(&req.Nodes, "nodes", 0, "with --gpu-class: the nodes one replica spans, over a fabric (vLLM only)")
	f.IntVar(&req.TP, "tp", 0, "the older name of --gpus")
	_ = f.MarkHidden("tp")
	f.BoolVar(&req.Fallbacks, "fallbacks", false, "also add the next best ways to run it, used when the first is busy")
	f.BoolVar(&detach, "detach", false, "start it and return: it runs on the control plane; `llmgw ctl jobs` follows it")
	f.StringVar(&req.Name, "name", "", "the catalog name (default: from the repository)")
	f.StringSliceVar(&req.Aliases, "alias", nil, "other names clients may use; repeatable")
	f.StringVar(&req.Description, "description", "", "what it is for, shown to clients in /v1/models")
	f.IntVar(&req.Replicas, "replicas", 1, "how many replicas to run")
	f.BoolVar(&onDemand, "on-demand", false, "run none until first asked for")
	f.IntVar(&req.CtxPool, "ctx-pool", 16384, "llama.cpp's --ctx-size, shared by its parallel slots")
	f.IntVar(&req.CPUCores, "cores", 0, "CPU cores a replica computes on; 0 takes half the smallest node it fits")
	f.IntVar(&req.ContextLength, "context-length", 4096, "what one request may use")
	return c
}

// followJob shows a job's steps as they happen, until it ends.
func followJob(cmd *cobra.Command, client func() *ctlClient, job modelserve.Job, doneSays string) error {
	out := cmd.OutOrStdout()
	fmt.Fprintf(out, "job %s\n", job.ID)
	shown := map[string]string{}
	for {
		for _, s := range job.Steps {
			line := s.State + " " + s.Detail
			if shown[s.Name] == line {
				continue
			}
			shown[s.Name] = line
			mark := map[string]string{modelserve.Running: "..", modelserve.Done: "ok",
				modelserve.Failed: "!!", modelserve.Skipped: "--"}[s.State]
			fmt.Fprintf(out, "%-3s %-11s %s\n", mark, s.Name, s.Detail)
		}
		switch job.State {
		case modelserve.Done:
			for _, st := range job.Steps {
				if st.Name == "placement" && st.State == modelserve.Skipped {
					fmt.Fprintf(out, "\n%s is in the catalog, waiting for room on the hardware it fits.\n", job.Model)
					return nil
				}
			}
			fmt.Fprintf(out, "\n%s %s.\n", job.Model, doneSays)
			return nil
		case modelserve.Failed:
			return errors.New(job.Error)
		}
		time.Sleep(2 * time.Second)
		if err := client().get(cmd.Context(), "/api/v1/jobs/"+job.ID, &job); err != nil {
			return err
		}
	}
}

// detached says where a job started without waiting for it can be followed.
func detached(cmd *cobra.Command, job modelserve.Job) error {
	fmt.Fprintf(cmd.OutOrStdout(), "started %s on the control plane\nfollow it: llmgw ctl jobs %s\n", job.ID, job.ID)
	return nil
}

func newJobsCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "jobs [JOB]",
		Short: "Models being served or updated, or one followed until it ends",
		Long: "Serving a model and moving one to new weights run as jobs on the control\n" +
			"plane, which carry on whether or not anything watches them; they are\n" +
			"kept until it restarts. With a job's id, its steps are shown as they\n" +
			"happen until it ends.",
		Args: cobra.MaximumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 1 {
				var job modelserve.Job
				if err := client().get(cmd.Context(), "/api/v1/jobs/"+args[0], &job); err != nil {
					return err
				}
				return followJob(cmd, client, job, "is done")
			}
			var jobs []modelserve.Job
			if err := client().get(cmd.Context(), "/api/v1/jobs", &jobs); err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(jobs)
			}
			if len(jobs) == 0 {
				fmt.Fprintln(cmd.OutOrStdout(), "no jobs since the control plane started")
				return nil
			}
			w := tabTo(cmd.OutOrStdout())
			fmt.Fprintln(w, "JOB\tMODEL\tSTATE\tSTARTED\tNOW")
			for _, j := range jobs {
				state := j.State
				for _, st := range j.Steps {
					if st.Name == "placement" && st.State == modelserve.Skipped {
						state = "waiting for room"
					}
				}
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", j.ID, or(j.Model, j.Request.Repo), state,
					ago(j.Started, time.Now()), truncate(jobNow(j), 80))
			}
			return w.Flush()
		},
	}
}