// Package eval asks a model a suite of questions and scores its answers, so
// new weights can be held to the quality of the old before they replace them.
//
// internal/bench measures how fast a replica answers; this measures whether
// the answers are any good. A repository that republishes its weights — a
// new quantisation, a fixed chat template — can make a model faster and
// worse at once, and nothing in the latency numbers says so. A suite is a
// handful of cases with known answers, scored by rules where a rule can
// decide (an exact string, a number, JSON of a shape) and by another model
// with a rubric where only judgement can.
package eval

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"path"
	"regexp"
	"strings"

	"go.yaml.in/yaml/v3"
)

// Suite is a set of cases a model is asked, and the score it must reach.
type Suite struct {
	// Name is what runs are filed under; the file's name without its
	// extension unless the suite gives one.
	Name        string `json:"name,omitempty" yaml:"name"`
	Description string `json:"description,omitempty" yaml:"description"`
	// Threshold is the score, 0 to 1, the suite is passed at. Zero is 0.8.
	Threshold float64 `json:"threshold,omitempty" yaml:"threshold"`
	// Judge is the model asked to grade the cases scored by a rubric, unless
	// a case names its own.
	Judge string `json:"judge,omitempty" yaml:"judge"`
	// System and MaxTokens are every case's unless it gives its own.
	System    string `json:"system,omitempty" yaml:"system"`
	MaxTokens int    `json:"max_tokens,omitempty" yaml:"max_tokens"`
	Cases     []Case `json:"cases" yaml:"cases"`

	// Hash is the sha256 of the text the suite was read from: two runs
	// under one name compare only when they asked the same questions.
	Hash string `json:"-" yaml:"-"`
}

// Message is one turn of a conversation a case sends.
type Message struct {
	Role    string `json:"role" yaml:"role"`
	Content string `json:"content" yaml:"content"`
}

// Case is one question and how its answer is scored: by exactly one of
// Exact, Contains, Regex, Number, JSONSchema or Judge.
type Case struct {
	// ID names the case in results; case-N, by its place, unless given.
	ID string `json:"id,omitempty" yaml:"id"`
	// Prompt is the user's message; Messages a whole conversation instead.
	Prompt   string    `json:"prompt,omitempty" yaml:"prompt"`
	Messages []Message `json:"messages,omitempty" yaml:"messages"`
	System   string    `json:"system,omitempty" yaml:"system"`
	// MaxTokens bounds the answer; the suite's, else 512.
	MaxTokens int `json:"max_tokens,omitempty" yaml:"max_tokens"`
	// Temperature is 0 unless given: a suite is a measurement, and sampling
	// would make two runs of the same weights disagree.
	Temperature *float64 `json:"temperature,omitempty" yaml:"temperature"`
	// Weight is how much the case counts against the others; 1 unless given.
	Weight float64 `json:"weight,omitempty" yaml:"weight"`

	// Exact is the whole answer, surrounding white space aside.
	Exact *string `json:"exact,omitempty" yaml:"exact"`
	// Contains are strings that must all be in the answer.
	Contains StringList `json:"contains,omitempty" yaml:"contains"`
	// Regex must match somewhere in the answer.
	Regex string `json:"regex,omitempty" yaml:"regex"`
	// IgnoreCase makes Exact, Contains and Regex blind to case.
	IgnoreCase bool `json:"ignore_case,omitempty" yaml:"ignore_case"`
	// Number is the value the last number in the answer must be, within
	// Tolerance either side.
	Number    *float64 `json:"number,omitempty" yaml:"number"`
	Tolerance float64  `json:"tolerance,omitempty" yaml:"tolerance"`
	// JSONSchema is a schema the answer, read as JSON, must satisfy.
	JSONSchema any `json:"json_schema,omitempty" yaml:"json_schema"`
	// Judge is a rubric another model grades the answer by, 0 to 1;
	// JudgeModel names that model, else the suite's judge.
	Judge      string `json:"judge,omitempty" yaml:"judge"`
	JudgeModel string `json:"judge_model,omitempty" yaml:"judge_model"`

	re     *regexp.Regexp
	schema *schema
}

// StringList is one string or a list of them.
type StringList []string

// UnmarshalYAML takes a string or a list.
func (l *StringList) UnmarshalYAML(n *yaml.Node) error {
	if n.Kind == yaml.ScalarNode {
		*l = StringList{n.Value}
		return nil
	}
	var v []string
	if err := n.Decode(&v); err != nil {
		return err
	}
	*l = v
	return nil
}

// UnmarshalJSON takes a string or a list.
func (l *StringList) UnmarshalJSON(b []byte) error {
	var s string
	if json.Unmarshal(b, &s) == nil {
		*l = StringList{s}
		return nil
	}
	var v []string
	if err := json.Unmarshal(b, &v); err != nil {
		return errors.New("contains is a string or a list of strings")
	}
	*l = v
	return nil
}

// Kind is how a case is scored: exact, contains, regex, number,
// json_schema or judge.
func (c Case) Kind() string {
	var kinds []string
	if c.Exact != nil {
		kinds = append(kinds, "exact")
	}
	if len(c.Contains) > 0 {
		kinds = append(kinds, "contains")
	}
	if c.Regex != "" {
		kinds = append(kinds, "regex")
	}
	if c.Number != nil {
		kinds = append(kinds, "number")
	}
	if c.JSONSchema != nil {
		kinds = append(kinds, "json_schema")
	}
	if c.Judge != "" {
		kinds = append(kinds, "judge")
	}
	return strings.Join(kinds, "+")
}

// DefaultThreshold is the score a suite that names none is passed at.
const DefaultThreshold = 0.8

// defaultMaxTokens bounds an answer when neither the case nor the suite do.
const defaultMaxTokens = 512

// Parse reads a suite: YAML (a mapping with its cases under cases:), or
// JSON lines (one case on each, and optionally a line without a prompt or
// messages that holds the suite's own settings). The name is the file's,
// used when the suite gives none.
func Parse(file string, data []byte) (Suite, error) {
	var s Suite
	var err error
	if isJSONLines(data) {
		s, err = parseJSONLines(data)
	} else {
		dec := yaml.NewDecoder(bytes.NewReader(data))
		dec.KnownFields(true)
		err = dec.Decode(&s)
		if errors.Is(err, io.EOF) {
			err = errors.New("the file is empty")
		}
	}
	if err != nil {
		return Suite{}, fmt.Errorf("%s: %w", file, err)
	}
	if s.Name == "" {
		base := path.Base(strings.ReplaceAll(file, "\\", "/"))
		s.Name = strings.TrimSuffix(base, path.Ext(base))
	}
	sum := sha256.Sum256(data)
	s.Hash = hex.EncodeToString(sum[:])
	if err := s.check(); err != nil {
		return Suite{}, fmt.Errorf("%s: %w", file, err)
	}
	return s, nil
}

// isJSONLines is whether every line that is not blank is a JSON object, and
// there is more than one, or the one there is is a case rather than a whole
// suite.
func isJSONLines(data []byte) bool {
	n := 0
	for _, line := range bytes.Split(data, []byte("\n")) {
		line = bytes.TrimSpace(line)
		if len(line) == 0 {
			continue
		}
		if line[0] != '{' || !json.Valid(line) {
			return false
		}
		n++
	}
	if n == 1 {
		var probe struct {
			Cases json.RawMessage `json:"cases"`
		}
		_ = json.Unmarshal(bytes.TrimSpace(data), &probe)
		return probe.Cases == nil
	}
	return n > 0
}

func parseJSONLines(data []byte) (Suite, error) {
	var s Suite
	for i, line := range bytes.Split(data, []byte("\n")) {
		line = bytes.TrimSpace(line)
		if len(line) == 0 {
			continue
		}
		var probe struct {
			Prompt   *string         `json:"prompt"`
			Messages json.RawMessage `json:"messages"`
		}
		_ = json.Unmarshal(line, &probe)
		dec := json.NewDecoder(bytes.NewReader(line))
		dec.DisallowUnknownFields()
		if probe.Prompt == nil && probe.Messages == nil {
			cases := s.Cases
			if err := dec.Decode(&s); err != nil {
				return Suite{}, fmt.Errorf("line %d: %w", i+1, err)
			}
			s.Cases = append(cases, s.Cases...)
			continue
		}
		var c Case
		if err := dec.Decode(&c); err != nil {
			return Suite{}, fmt.Errorf("line %d: %w", i+1, err)
		}
		s.Cases = append(s.Cases, c)
	}
	return s, nil
}

// check fills in what a suite left out and refuses what cannot be run.
func (s *Suite) check() error {
	if len(s.Cases) == 0 {
		return errors.New("no cases: a suite has them under cases:, or one on each line of JSON")
	}
	if s.Threshold < 0 || s.Threshold > 1 {
		return fmt.Errorf("threshold %v is not 0 to 1", s.Threshold)
	}
	if s.Threshold == 0 {
		s.Threshold = DefaultThreshold
	}
	seen := map[string]bool{}
	for i := range s.Cases {
		c := &s.Cases[i]
		if c.ID == "" {
			c.ID = fmt.Sprintf("case-%d", i+1)
		}
		if seen[c.ID] {
			return fmt.Errorf("two cases are called %s", c.ID)
		}
		seen[c.ID] = true
		if err := c.check(); err != nil {
			return fmt.Errorf("case %s: %w", c.ID, err)
		}
	}
	return nil
}

func (c *Case) check() error {
	if strings.TrimSpace(c.Prompt) == "" && len(c.Messages) == 0 {
		return errors.New("no prompt and no messages: nothing to ask")
	}
	if c.Prompt != "" && len(c.Messages) > 0 {
		return errors.New("a prompt and messages both: give one")
	}
	for _, m := range c.Messages {
		switch m.Role {
		case "system", "user", "assistant":
		default:
			return fmt.Errorf("a message's role is system, user or assistant, not %q", m.Role)
		}
	}
	switch k := c.Kind(); {
	case k == "":
		return errors.New("no way to score it: give one of exact, contains, regex, number, json_schema or judge")
	case strings.Contains(k, "+"):
		return fmt.Errorf("scored by %s: give one", strings.ReplaceAll(k, "+", " and "))
	}
	if c.Weight < 0 {
		return fmt.Errorf("weight %v is negative", c.Weight)
	}
	if c.Weight == 0 {
		c.Weight = 1
	}
	if c.Tolerance < 0 {
		return fmt.Errorf("tolerance %v is negative", c.Tolerance)
	}
	if c.Regex != "" {
		expr := c.Regex
		if c.IgnoreCase {
			expr = "(?i)" + expr
		}
		re, err := regexp.Compile(expr)
		if err != nil {
			return fmt.Errorf("the regex: %w", err)
		}
		c.re = re
	}
	if c.JSONSchema != nil {
		sc, err := compileSchema(c.JSONSchema)
		if err != nil {
			return fmt.Errorf("the JSON schema: %w", err)
		}
		c.schema = sc
	}
	return nil
}

// NeedsJudge is whether any case is graded by a model.
func (s Suite) NeedsJudge() bool {
	for _, c := range s.Cases {
		if c.Judge != "" {
			return true
		}
	}
	return false
}