package cli

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/placement"
	"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
)

// Machines of the judges' own. Under load on a 4-core control plane that was
// also a node, 38 of 60 requests were refused for want of a judgment; with the
// judges on a node of their own the same test refused 2. A judge machine is
// added as a node carrying judgeRoleLabel, and the judge's profile selects
// that label, so nothing else is placed beside it and it on nothing else.

// judgeRoleLabel marks a node as the judges'.
const (
	judgeRoleLabel = placement.RoleLabel
	judgeRoleValue = "judge"
)

// judgeMachine is a machine for the judges: user@host, and the connection the
// setup screen already made to it, when it made one.
type judgeMachine struct {
	target string
	remote onboard.Remote
}

// ensureJudgeNodes adds each judge machine as a node labelled for the
// judges, and says how many cores the smallest of them has.
func (in *serverInit) ensureJudgeNodes(ctx context.Context) error {
	if len(in.judgeMachines) == 0 {
		return nil
	}
	na := nodeAdd{serverCfg: in.configPath(), plain: true}
	na.opts.PKI = in.pkiDir
	na.opts.NoFirewall = in.noFirewall
	if err := na.prepare(); err != nil {
		return err
	}
	in.judgeCPUs = nil
	for _, jm := range in.judgeMachines {
		o := na.target(jm.target)
		fmt.Fprintf(in.out, "  ..    adding %s as a node for the judges\n", o.Addr)
		r := onboard.New(o, nil)
		if jm.remote != nil {
			r.Remote = jm.remote
		}
		res, err := r.Run(ctx)
		if err != nil {
			return fmt.Errorf("%s: %w", o.Addr, err)
		}
		if err := na.watch(ctx, res.Node, func(onboard.Event) {}); err != nil && err != errNoAdmin {
			return fmt.Errorf("%s: %w", res.Node, err)
		}
		if out, err := r.Remote.Run(ctx, "nproc", nil); err == nil {
			if n, err := strconv.Atoi(strings.TrimSpace(string(out))); err == nil {
				in.judgeCPUs = append(in.judgeCPUs, n)
			}
		}
		if _, err := changeStoredSettings(in.configPath(), "a node for the judges, from init", func(doc *settings.Doc) error {
			return doc.SetYAML([]string{"agents", "nodes", res.Node, "labels", judgeRoleLabel}, judgeRoleValue)
		}); err != nil {
			return fmt.Errorf("labelling %s for the judges: %w", res.Node, err)
		}
		in.done("%s is a node for the judges alone (%s: %s)", res.Node, judgeRoleLabel, judgeRoleValue)
	}
	return nil
}

// judgePlacement is the judge's cores and where it may run: on the judge
// machines when there are some, by init's rule over their cores; otherwise
// on any node, with this machine's plan.
func (in *serverInit) judgePlacement() (cores int, selector map[string]string) {
	if len(in.judgeMachines) == 0 {
		return in.cpuPlan().judge, nil
	}
	return refitCores(in.judgeFor().Replicas, in.judgeCPUs), map[string]string{judgeRoleLabel: judgeRoleValue}
}

// judgeMachineFlag is --judge-node, given again for each machine.
type judgeMachineFlag struct{ to *[]judgeMachine }

func (f judgeMachineFlag) String() string {
	if f.to == nil {
		return ""
	}
	var out []string
	for _, m := range *f.to {
		out = append(out, m.target)
	}
	return strings.Join(out, ",")
}

func (f judgeMachineFlag) Set(v string) error {
	if strings.TrimSpace(v) == "" {
		return fmt.Errorf("--judge-node wants user@host")
	}
	*f.to = append(*f.to, judgeMachine{target: strings.TrimSpace(v)})
	return nil
}

func (judgeMachineFlag) Type() string { return "user@host" }

// judgeReadyWait is how long init waits for its judge to load before timing
// it: the weights come from the mirror, and llama.cpp loads in seconds.
const judgeReadyWait = 15 * time.Minute

// checkJudge times the judge once it is up, and keeps what it found for the
// closing text: a machine too slow for its judge is said before its first
// refused request, not after. The budget is not changed without asking.
func (in *serverInit) checkJudge(ctx context.Context) {
	c := in.adminAPI()
	deadline := time.Now().Add(judgeReadyWait)
	for {
		var judges []adminapi.JudgeView
		if err := c.get(ctx, "/api/v1/guard/judges", &judges); err == nil {
			ready := false
			for _, j := range judges {
				ready = ready || (j.Model == guardModelName && j.Ready > 0)
			}
			if ready {
				break
			}
		}
		if time.Now().After(deadline) || ctx.Err() != nil {
			in.judgeNote = "  !!    the judge was not up to be timed; once it is: llmgw ctl guard judge check\n"
			return
		}
		time.Sleep(10 * time.Second)
	}
	fmt.Fprintf(in.out, "  ..    timing the judge\n")
	v, err := checkJudge(ctx, c, guardModelName, 0)
	if err != nil {
		in.judgeNote = "  !!    the judge could not be timed (" + err.Error() + "); later: llmgw ctl guard judge check\n"
		return
	}
	if v.Verdict == "ok" {
		in.done("the judge reads a prompt in %.1f s, of the %.1f s a request waits; about %.1f requests a second",
			float64(v.SlowestMs)/1000, float64(v.BudgetMs)/1000, v.PerSecond/2)
		return
	}
	in.judgeNote = "The judge is " + v.Verdict + " on this cluster:\n" + indent(judgeCheckText(v), "  ") + "\n"
}