package eval

import (
	"encoding/json"
	"fmt"
	"math"
	"regexp"
	"strconv"
	"strings"
)

// verdict is how one answer scored, 0 to 1, and why.
type verdict struct {
	score  float64
	detail string
}

func pass(detail string) verdict { return verdict{1, detail} }
func fail(detail string) verdict { return verdict{0, detail} }

var thinkBlock = regexp.MustCompile(`(?s)<think>.*?</think>`)

// cleanAnswer is what a model said, without a reasoning model's thinking —
// what a reader is shown — and without the white space around it.
func cleanAnswer(s string) string {
	s = thinkBlock.ReplaceAllString(s, "")
	// A reply cut off while still thinking has an opening tag and no close.
	if i := strings.Index(s, "<think>"); i >= 0 {
		s = s[:i]
	}
	return strings.TrimSpace(s)
}

// scoreRule scores an answer by a case's rule; a case graded by a judge is
// scored elsewhere, since it takes another request.
func scoreRule(c *Case, answer string) verdict {
	fold := func(s string) string {
		if c.IgnoreCase {
			return strings.ToLower(s)
		}
		return s
	}
	switch {
	case c.Exact != nil:
		if fold(answer) == fold(strings.TrimSpace(*c.Exact)) {
			return pass("the answer is exactly that")
		}
		return fail(fmt.Sprintf("wanted exactly %q", strings.TrimSpace(*c.Exact)))
	case len(c.Contains) > 0:
		var missing []string
		for _, want := range c.Contains {
			if !strings.Contains(fold(answer), fold(want)) {
				missing = append(missing, strconv.Quote(want))
			}
		}
		if len(missing) == 0 {
			return pass("has everything asked for")
		}
		return fail("has no " + strings.Join(missing, ", "))
	case c.re != nil:
		if c.re.MatchString(answer) {
			return pass("matches")
		}
		return fail("does not match " + c.re.String())
	case c.Number != nil:
		got, ok := lastNumber(answer)
		if !ok {
			return fail(fmt.Sprintf("no number in the answer; wanted %v", *c.Number))
		}
		if math.Abs(got-*c.Number) <= c.Tolerance {
			return pass(fmt.Sprintf("read %v", got))
		}
		want := fmt.Sprintf("%v", *c.Number)
		if c.Tolerance > 0 {
			want += fmt.Sprintf(" ± %v", c.Tolerance)
		}
		return fail(fmt.Sprintf("read %v, wanted %s", got, want))
	case c.schema != nil:
		v, err := readJSON(answer)
		if err != nil {
			return fail(err.Error())
		}
		if why := c.schema.validate(v, ""); why != "" {
			return fail(why)
		}
		return pass("JSON of the shape asked for")
	}
	return fail("no rule to score it by")
}

// number is a number as a model writes one: a sign, thousands separated by
// commas, a decimal part, an exponent.
var number = regexp.MustCompile(`[-+]?(?:\d{1,3}(?:,\d{3})+|\d+)(?:\.\d+)?(?:[eE][-+]?\d+)?`)

// lastNumber is the last number in a text: a model that shows its working
// ends with the result, and one that does not says only the result.
func lastNumber(s string) (float64, bool) {
	all := number.FindAllStringIndex(s, -1)
	if len(all) == 0 {
		return 0, false
	}
	i, j := all[len(all)-1][0], all[len(all)-1][1]
	// A hyphen after a digit or letter joins (3-5, 2026-10-04); it is not a
	// minus sign.
	if (s[i] == '-' || s[i] == '+') && i > 0 && isWordByte(s[i-1]) {
		i++
	}
	f, err := strconv.ParseFloat(strings.ReplaceAll(s[i:j], ",", ""), 64)
	return f, err == nil
}

func isWordByte(b byte) bool {
	return b >= '0' && b <= '9' || b >= 'a' && b <= 'z' || b >= 'A' && b <= 'Z'
}

var fence = regexp.MustCompile("(?s)```(?:json|JSON)?\\s*(.*?)```")

// readJSON reads an answer as JSON: the whole of it, else what is inside a
// code fence, else from the first brace or bracket to the last — models
// asked for JSON wrap it in a sentence often enough that refusing that would
// measure manners rather than the shape.
func readJSON(answer string) (any, error) {
	try := []string{answer}
	if m := fence.FindStringSubmatch(answer); m != nil {
		try = append(try, m[1])
	}
	for _, pair := range [][2]string{{"{", "}"}, {"[", "]"}} {
		i, j := strings.Index(answer, pair[0]), strings.LastIndex(answer, pair[1])
		if i >= 0 && j > i {
			try = append(try, answer[i:j+1])
		}
	}
	for _, t := range try {
		var v any
		if json.Unmarshal([]byte(strings.TrimSpace(t)), &v) == nil {
			return v, nil
		}
	}
	if answer == "" {
		return nil, fmt.Errorf("the answer is empty, not JSON")
	}
	return nil, fmt.Errorf("the answer is not JSON: %s", clip(answer, 80))
}

func clip(s string, n int) string {
	s = strings.Join(strings.Fields(s), " ")
	if len(s) <= n {
		return s
	}
	return s[:n-3] + "..."
}

// judgePrompt asks a model to grade an answer by a rubric.
const judgePrompt = "You grade an answer to a question against a rubric. Read the question, the answer " +
	"and the rubric, then reply with only a JSON object: {\"score\": a number from 0 to 1, " +
	"\"reason\": \"one short sentence\"}. 1 means the answer meets the rubric fully, 0 not at all."

// judgeRequest is what a judge is asked about one answer.
func judgeRequest(c *Case, answer string) []Message {
	var q strings.Builder
	for _, m := range askedOf(c, "") {
		if m.Role == "system" {
			continue
		}
		fmt.Fprintf(&q, "%s: %s\n", m.Role, m.Content)
	}
	return []Message{
		{Role: "system", Content: judgePrompt},
		{Role: "user", Content: fmt.Sprintf("Question:\n%s\nAnswer:\n%s\n\nRubric:\n%s", q.String(), answer, c.Judge)},
	}
}

var (
	passWord = regexp.MustCompile(`(?i)\b(pass|passed|yes|correct)\b`)
	failWord = regexp.MustCompile(`(?i)\b(fail|failed|no|incorrect|wrong)\b`)
)

// readJudgement reads a judge's reply: a JSON object with a score (or a
// pass true or false), else a lone number, else the words pass or fail.
func readJudgement(reply string) (verdict, error) {
	reply = cleanAnswer(reply)
	if v, err := readJSON(reply); err == nil {
		if m, ok := v.(map[string]any); ok {
			reason, _ := m["reason"].(string)
			if f, ok := m["score"].(float64); ok {
				return verdict{clamp(f), "the judge: " + clip(reason, 120)}, nil
			}
			if s, ok := m["score"].(string); ok {
				if f, err := strconv.ParseFloat(strings.TrimSpace(s), 64); err == nil {
					return verdict{clamp(f), "the judge: " + clip(reason, 120)}, nil
				}
			}
			if b, ok := m["pass"].(bool); ok {
				if b {
					return pass("the judge: " + clip(reason, 120)), nil
				}
				return fail("the judge: " + clip(reason, 120)), nil
			}
		}
	}
	if f, err := strconv.ParseFloat(strings.TrimSpace(reply), 64); err == nil {
		return verdict{clamp(f), "the judge gave " + strings.TrimSpace(reply)}, nil
	}
	p, f := passWord.MatchString(reply), failWord.MatchString(reply)
	switch {
	case p && !f:
		return pass("the judge: " + clip(reply, 120)), nil
	case f && !p:
		return fail("the judge: " + clip(reply, 120)), nil
	}
	return verdict{}, fmt.Errorf("the judge's answer could not be read as a score: %s", clip(reply, 120))
}

func clamp(f float64) float64 { return math.Max(0, math.Min(1, f)) }