package cli

import (
	"fmt"
	"slices"
	"strings"

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

// knownJudge is a judge model llmgw knows where to fetch and how to serve:
// what a policy may name without anyone saying where its weights are.
type knownJudge struct {
	Name, Repo, File, Format string
	// GPU is a judge too slow on CPUs to wait for: served on one card.
	GPU   bool
	About string
}

// knownJudges, the default first. Measured on the reference cluster against
// ordinary coding requests ("continue", "kill the process on port 8080") and
// harmful ones: Llama Guard 3 1B scored both alike, about one in two
// ordinary ones refused; Qwen3Guard-Gen 0.6B kept them apart, 0.15 against
// 1.00 at the worst, at about two seconds a check on two CPU threads.
var knownJudges = []knownJudge{
	{Name: "qwen3guard-0.6b", Repo: "mradermacher/Qwen3Guard-Gen-0.6B-GGUF", File: "Qwen3Guard-Gen-0.6B.Q8_0.gguf",
		Format: "qwen3guard", About: "Qwen3Guard-Gen 0.6B on CPUs, 0.8 GB: the default"},
	{Name: "qwen3guard-4b", Repo: "mradermacher/Qwen3Guard-Gen-4B-GGUF", File: "Qwen3Guard-Gen-4B.Q4_K_M.gguf",
		Format: "qwen3guard", GPU: true, About: "Qwen3Guard-Gen 4B on a GPU, 2.5 GB"},
	{Name: "qwen3guard-8b", Repo: "mradermacher/Qwen3Guard-Gen-8B-GGUF", File: "Qwen3Guard-Gen-8B.Q4_K_M.gguf",
		Format: "qwen3guard", GPU: true, About: "Qwen3Guard-Gen 8B on a GPU, 4.7 GB"},
	{Name: "llama-guard-3-1b", Repo: "QuantFactory/Llama-Guard-3-1B-GGUF", File: "Llama-Guard-3-1B.Q4_K_M.gguf",
		Format: "llamaguard3", About: "Llama Guard 3 1B on CPUs, 0.9 GB: refuses ordinary coding requests"},
	{Name: "llama-guard-3-8b", Repo: "QuantFactory/Llama-Guard-3-8B-GGUF", File: "Llama-Guard-3-8B.Q4_K_M.gguf",
		Format: "llamaguard3", GPU: true, About: "Llama Guard 3 8B on a GPU, 4.6 GB"},
}

const (
	defaultJudge    = "qwen3guard-0.6b"
	defaultGPUJudge = "qwen3guard-4b"
)

// baselineCategories are what the baseline policy counts of each format's
// verdicts: not PII, which is the patterns', nor what refuses ordinary work
// — copyright, unethical acts, politically sensitive topics; Llama Guard's
// defamation, specialized advice, intellectual property and elections.
var baselineCategories = map[string][]string{
	"qwen3guard":  {"Violent", "Non-violent Illegal Acts", "Sexual Content or Sexual Acts", "Suicide & Self-Harm", "Jailbreak"},
	"llamaguard3": {"S1", "S2", "S3", "S4", "S7", "S9", "S10", "S11", "S12"},
}

// reported is whether a judge of that format ever names category c: its own
// names, or for Llama Guard its codes.
func reported(format, c string) bool {
	if slices.ContainsFunc(guard.DefaultCategories(format), func(n string) bool { return strings.EqualFold(n, c) }) {
		return true
	}
	var n int
	_, err := fmt.Sscanf(c, "S%d", &n)
	return format == "llamaguard3" && err == nil && n >= 1 && n <= 14
}

// categoriesFor is what a policy asking these judges should count, when the
// categories it has are ones none of them names: a policy moved from Llama
// Guard's S1 to Qwen3Guard, which says Violent, would otherwise count
// nothing. Nil keeps what it has; so does a judge llmgw does not know.
func categoriesFor(judges, have []string) []string {
	var formats []string
	for _, j := range judges {
		k, ok := knownJudgeNamed(j)
		if !ok {
			return nil
		}
		if !slices.Contains(formats, k.Format) {
			formats = append(formats, k.Format)
		}
	}
	if len(have) == 0 {
		return nil
	}
	stray := false
	for _, f := range formats {
		for _, c := range have {
			if !reported(f, c) && !slices.ContainsFunc(formats, func(o string) bool { return reported(o, c) }) {
				stray = true
			}
		}
		// Every judge must have something it counts.
		if !slices.ContainsFunc(have, func(c string) bool { return reported(f, c) }) {
			stray = true
		}
	}
	if !stray {
		return nil
	}
	var out []string
	for _, f := range formats {
		out = append(out, baselineCategories[f]...)
	}
	return out
}

func knownJudgeNamed(name string) (knownJudge, bool) {
	i := slices.IndexFunc(knownJudges, func(k knownJudge) bool { return k.Name == name })
	if i < 0 {
		return knownJudge{}, false
	}
	return knownJudges[i], true
}

func knownJudgeNames() string {
	var n []string
	for _, k := range knownJudges {
		n = append(n, k.Name)
	}
	return strings.Join(n, ", ")
}

// judgeServeOptions are how a judge is served beyond which weights.
type judgeServeOptions struct {
	Replicas, Cores int
	// GPUClass is the cards a GPU judge runs on; empty lets the plan choose.
	GPUClass string
}

// serveRequest is the serve job for a judge. A judge on CPUs is critical,
// as every judged request waits on it; one on a GPU is standard, so it takes
// a card only from what may give one up, and the judge on CPUs after it in
// the policies stands in while it has none.
func (k knownJudge) serveRequest(o judgeServeOptions) modelserve.Request {
	req := modelserve.Request{
		Repo: k.Repo, File: k.File, Name: k.Name, GuardFormat: k.Format,
		Description: "The guardrails' judge: reads prompts and answers for unsafe content. Not offered to clients.",
		Replicas:    o.Replicas, CtxPool: 8192, ContextLength: 4096,
	}
	if k.GPU {
		req.GPUClass, req.GPUsPerNode, req.Class = o.GPUClass, 1, catalog.ClassStandard
		if req.Replicas == 0 {
			req.Replicas = 1
		}
		return req
	}
	req.CPU, req.CPUCores = true, o.Cores
	if req.Replicas == 0 {
		req.Replicas = 2
	}
	return req
}

func (k knownJudge) String() string {
	where := "CPUs"
	if k.GPU {
		where = "a GPU"
	}
	return fmt.Sprintf("%s (%s, %s/%s, on %s)", k.Name, k.Format, k.Repo, k.File, where)
}