package eval

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// Caller sends a chat completion to a model as the gateway does for itself:
// the router's Replay, past admission, accounting and the guardrails — a
// suite measures the weights, and a tenant's quota or a guardrail's refusal
// would measure something else. It answers the status, the body, and how
// long it took.
type Caller func(ctx context.Context, path, model string, body []byte) (int, []byte, time.Duration, error)

// Options are how a suite is run.
type Options struct {
	// Model is the catalog model asked.
	Model string
	// Concurrency is how many cases are asked at once; 4 unless given.
	Concurrency int
	// Judge grades the cases scored by a rubric, instead of the suite's.
	Judge string
	// MinScore is the score to pass at, instead of the suite's threshold.
	MinScore float64
	// CaseTimeout bounds one case, its judging included; five minutes
	// unless given, which a model loading on demand needs.
	CaseTimeout time.Duration
	// Progress is told as each case finishes.
	Progress func(done, total int)
}

// PassAt is the score a case passes at: a judge's 0.5 is a coin toss, and
// counts as passed so a lenient rubric is not a hidden zero.
const PassAt = 0.5

// keptAnswer is how much of an answer is kept with its result.
const keptAnswer = 2000

// Evaluate asks a model every case of a suite and scores the answers. A
// case that got no answer scores 0 with its error; a suite that cannot be
// run at all — a rubric and no judge to read it — is an error before
// anything is asked.
func Evaluate(ctx context.Context, call Caller, s Suite, o Options) (store.EvalRun, error) {
	if o.Model == "" {
		return store.EvalRun{}, errors.New("no model to ask")
	}
	judge := o.Judge
	if judge == "" {
		judge = s.Judge
	}
	for _, c := range s.Cases {
		if c.Judge != "" && c.JudgeModel == "" && judge == "" {
			return store.EvalRun{}, fmt.Errorf("case %s is graded by a judge and the suite names none: "+
				"give one with --judge MODEL, or judge: in the suite", c.ID)
		}
	}
	threshold := s.Threshold
	if o.MinScore > 0 {
		threshold = o.MinScore
	}
	if threshold <= 0 {
		threshold = DefaultThreshold
	}
	workers := o.Concurrency
	if workers <= 0 {
		workers = 4
	}
	timeout := o.CaseTimeout
	if timeout <= 0 {
		timeout = 5 * time.Minute
	}

	run := store.EvalRun{Model: o.Model, Suite: s.Name, SuiteHash: s.Hash, At: time.Now().UTC(),
		Threshold: threshold, Cases: len(s.Cases), Results: make([]store.EvalCase, len(s.Cases))}
	if s.NeedsJudge() {
		run.Judge = judge
	}
	todo := make(chan int)
	var wg sync.WaitGroup
	var mu sync.Mutex
	done := 0
	for range min(workers, len(s.Cases)) {
		wg.Go(func() {
			for i := range todo {
				cctx, cancel := context.WithTimeout(ctx, timeout)
				res := askOne(cctx, call, s, &s.Cases[i], o.Model, judge)
				cancel()
				mu.Lock()
				run.Results[i] = res
				done++
				if o.Progress != nil {
					o.Progress(done, len(s.Cases))
				}
				mu.Unlock()
			}
		})
	}
	for i := range s.Cases {
		select {
		case todo <- i:
		case <-ctx.Done():
		}
		if ctx.Err() != nil {
			break
		}
	}
	close(todo)
	wg.Wait()
	run.TookMs = time.Since(run.At).Milliseconds()

	var sum, weights float64
	firstErr := ""
	for i, r := range run.Results {
		if r.ID == "" {
			// Not asked: the run was stopped first.
			r = store.EvalCase{ID: s.Cases[i].ID, Kind: s.Cases[i].Kind(), Weight: s.Cases[i].Weight,
				Error: "not asked: the run was stopped"}
			run.Results[i] = r
		}
		sum += r.Weight * r.Score
		weights += r.Weight
		if !r.Passed {
			run.Below++
		}
		if r.Error != "" {
			run.Errors++
			if firstErr == "" {
				firstErr = r.Error
			}
		}
	}
	if weights > 0 {
		run.Score = sum / weights
	}
	run.Passed = run.Score >= threshold
	if run.Errors == run.Cases {
		run.Passed = false
		run.Error = "no case was answered: " + firstErr
	}
	return run, nil
}

// askedOf is the conversation a case sends.
func askedOf(c *Case, suiteSystem string) []Message {
	var out []Message
	system := c.System
	if system == "" {
		system = suiteSystem
	}
	if system != "" && (len(c.Messages) == 0 || c.Messages[0].Role != "system") {
		out = append(out, Message{Role: "system", Content: system})
	}
	if len(c.Messages) > 0 {
		return append(out, c.Messages...)
	}
	return append(out, Message{Role: "user", Content: c.Prompt})
}

// askOne asks one case and scores its answer.
func askOne(ctx context.Context, call Caller, s Suite, c *Case, model, judge string) store.EvalCase {
	res := store.EvalCase{ID: c.ID, Kind: c.Kind(), Weight: c.Weight}
	start := time.Now()
	defer func() { res.TookMs = time.Since(start).Milliseconds() }()

	maxTokens := c.MaxTokens
	if maxTokens <= 0 {
		maxTokens = s.MaxTokens
	}
	if maxTokens <= 0 {
		maxTokens = defaultMaxTokens
	}
	temp := 0.0
	if c.Temperature != nil {
		temp = *c.Temperature
	}
	answer, finish, err := chat(ctx, call, model, askedOf(c, s.System), maxTokens, temp)
	if err != nil {
		res.Error = err.Error()
		return res
	}
	answer = cleanAnswer(answer)
	res.Answer = answer
	if len(res.Answer) > keptAnswer {
		res.Answer = res.Answer[:keptAnswer]
	}

	var v verdict
	if c.Judge != "" {
		jm := c.JudgeModel
		if jm == "" {
			jm = judge
		}
		reply, _, err := chat(ctx, call, jm, judgeRequest(c, answer), 256, 0)
		if err != nil {
			res.Error = "the judge " + jm + ": " + err.Error()
			return res
		}
		if v, err = readJudgement(reply); err != nil {
			res.Error = err.Error()
			return res
		}
	} else {
		v = scoreRule(c, answer)
	}
	res.Score, res.Detail = v.score, v.detail
	res.Passed = res.Score >= PassAt
	if finish == "length" && !res.Passed {
		res.Detail += fmt.Sprintf("; the answer was cut off at %d tokens", maxTokens)
	}
	if answer == "" && !res.Passed {
		res.Detail += "; it answered nothing"
	}
	return res
}

// chat asks a model one conversation and says what it answered and why it
// stopped.
func chat(ctx context.Context, call Caller, model string, msgs []Message, maxTokens int, temp float64) (string, string, error) {
	body, err := json.Marshal(map[string]any{"model": model, "messages": msgs, "max_tokens": maxTokens,
		"temperature": temp, "stream": false})
	if err != nil {
		return "", "", err
	}
	status, out, _, err := call(ctx, "/v1/chat/completions", model, body)
	if err != nil {
		return "", "", err
	}
	if status != http.StatusOK {
		var e struct {
			Error struct {
				Message string `json:"message"`
			} `json:"error"`
		}
		msg := strings.TrimSpace(string(out))
		if json.Unmarshal(out, &e) == nil && e.Error.Message != "" {
			msg = e.Error.Message
		}
		return "", "", fmt.Errorf("%s answered %d: %s", model, status, clip(msg, 200))
	}
	var resp struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(out, &resp); err != nil || len(resp.Choices) == 0 {
		return "", "", fmt.Errorf("%s's answer is not a chat completion: %s", model, clip(string(out), 120))
	}
	return resp.Choices[0].Message.Content, resp.Choices[0].FinishReason, nil
}