package cli

import (
	"encoding/json"
	"fmt"
	"slices"
	"text/tabwriter"

	"github.com/spf13/cobra"

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

func newModelCanaryCmd(client func() *ctlClient) *cobra.Command {
	var (
		cn            catalog.Canary
		abort, detach bool
	)
	c := &cobra.Command{
		Use:   "canary MODEL",
		Short: "Try a model's newer weights on a share of its requests before moving to them",
		Long: "The newer files of the model's repository are fetched and run beside it\n" +
			"as MODEL-canary, one replica, and given --percent of MODEL's requests. A\n" +
			"session (x-session-id, or the user field) stays on one side; the\n" +
			"x-llmgw-variant header (canary or stable) pins a request to one. Clients\n" +
			"cannot name the canary itself.\n\n" +
			"The two sides are compared as they serve. Once the canary has served\n" +
			"--min-requests, it is taken out by itself if its error rate is\n" +
			"--max-error-increase above the model's, or its time to the first token\n" +
			"(or, unstreamed, to the answer) --max-ttft-ratio times the model's, and\n" +
			"canary_rolled_back goes to the notification webhooks.\n\n" +
			"Run again with --percent to change its share, --abort to take it out;\n" +
			"`llmgw ctl model promote MODEL` moves the model to it.",
		Example: "  llmgw ctl model canary qwen-7b --percent 10\n" +
			"  llmgw ctl model canary qwen-7b --percent 50\n" +
			"  llmgw ctl model canary qwen-7b --abort",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			name := args[0]
			var out map[string]string
			if abort {
				if err := client().delete(cmd.Context(), "/api/v1/models/"+name+"/canary", &out); err != nil {
					return err
				}
				fmt.Fprintln(cmd.OutOrStdout(), out["result"])
				return nil
			}
			var have []modelserve.CanaryView
			if err := client().get(cmd.Context(), "/api/v1/canaries", &have); err != nil {
				return err
			}
			if slices.ContainsFunc(have, func(v modelserve.CanaryView) bool { return v.Model == name }) {
				if !cmd.Flags().Changed("percent") {
					return fmt.Errorf("%s already has a canary: give --percent to change its share, or --abort", name)
				}
				body, _ := json.Marshal(map[string]int{"percent": cn.Percent})
				if err := client().put(cmd.Context(), "/api/v1/models/"+name+"/canary", body, &out); err != nil {
					return err
				}
				fmt.Fprintln(cmd.OutOrStdout(), out["result"])
				return nil
			}
			body, _ := json.Marshal(cn)
			var job modelserve.Job
			if err := client().post(cmd.Context(), "/api/v1/models/"+name+"/canary", body, &job); err != nil {
				return err
			}
			if detach {
				return detached(cmd, job)
			}
			return followJob(cmd, client, job, "has a canary")
		},
	}
	f := c.Flags()
	f.IntVar(&cn.Percent, "percent", 10, "share of the model's requests the canary is given, 0 to 100")
	f.IntVar(&cn.MinRequests, "min-requests", 50, "requests the canary serves before it is judged")
	f.Float64Var(&cn.MaxErrorIncrease, "max-error-increase", 0.05, "error rate above the model's that rolls it back (0.05 is five points)")
	f.Float64Var(&cn.MaxTTFTRatio, "max-ttft-ratio", 1.5, "times the model's time to the first token that rolls it back")
	f.BoolVar(&abort, "abort", false, "take the canary out; the model serves all its requests again")
	f.BoolVar(&detach, "detach", false, "start it and return: it runs on the control plane; `llmgw ctl jobs` follows it")
	return c
}

func newModelPromoteCmd(client func() *ctlClient) *cobra.Command {
	var detach bool
	c := &cobra.Command{
		Use:   "promote MODEL",
		Short: "Move a model to its canary's weights, and take the canary out",
		Long: "The update the canary was a trial of: the catalog names the canary's\n" +
			"files, the model's replicas are replaced one at a time, and the canary\n" +
			"is removed. `llmgw ctl model rollback MODEL` goes back.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var job modelserve.Job
			if err := client().post(cmd.Context(), "/api/v1/models/"+args[0]+"/promote", []byte("{}"), &job); err != nil {
				return err
			}
			if detach {
				return detached(cmd, job)
			}
			return followJob(cmd, client, job, "runs the canary's weights")
		},
	}
	c.Flags().BoolVar(&detach, "detach", false, "start it and return: it runs on the control plane; `llmgw ctl jobs` follows it")
	return c
}

func newModelCanariesCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "canaries",
		Short: "Models trying a canary, and how each side has served since it began",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			var vs []modelserve.CanaryView
			if err := client().get(cmd.Context(), "/api/v1/canaries", &vs); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(vs)
			}
			if len(vs) == 0 {
				fmt.Fprintln(cmd.OutOrStdout(), "No model has a canary.")
				return nil
			}
			w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0)
			fmt.Fprintln(w, "MODEL\tCANARY\tSHARE\tREQUESTS\tERRORS\tFIRST TOKEN\tANSWER")
			for _, v := range vs {
				side := func(label string, req int64, errRate, ttft, lat float64) {
					fmt.Fprintf(w, "%s\t%s\t\t%d\t%.1f%%\t%s\t%s\n", "", label, req, 100*errRate, meanMs(ttft), meanMs(lat))
				}
				fmt.Fprintf(w, "%s\t%s\t%d%%\t\t\t\t\n", v.Model, v.Canary.Model, v.Canary.Percent)
				if s := v.Serving; s != nil {
					side("  canary", s.Candidate.Requests, s.Candidate.ErrorRate(), s.Candidate.TTFTMs, s.Candidate.LatencyMs)
					side("  stable", s.Stable.Requests, s.Stable.ErrorRate(), s.Stable.TTFTMs, s.Stable.LatencyMs)
				}
			}
			return w.Flush()
		},
	}
}

func meanMs(v float64) string {
	if v <= 0 {
		return "-"
	}
	return fmt.Sprintf("%.0f ms", v)
}