package eval

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"testing/fstest"
	"time"

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

const arith = `
name: arith
threshold: 0.75
judge: grader
cases:
  - id: add
    prompt: "What is 17 + 25? Answer with the number only."
    number: 42
  - id: mul
    prompt: "What is 6 * 7?"
    number: 42
    weight: 2
  - id: capital
    prompt: "What is the capital of France?"
    contains: Paris
  - id: polite
    prompt: "Say hello."
    judge: "A greeting of any kind passes."
`

// fakeModel answers by the prompt, and as a judge by what it was asked to
// grade; it counts what it was asked.
type fakeModel struct {
	mu             sync.Mutex
	answers        map[string]string
	asked          []string
	inFlight, most atomic.Int32
}

func (f *fakeModel) call(_ context.Context, path, model string, body []byte) (int, []byte, time.Duration, error) {
	n := f.inFlight.Add(1)
	defer f.inFlight.Add(-1)
	for {
		m := f.most.Load()
		if n <= m || f.most.CompareAndSwap(m, n) {
			break
		}
	}
	time.Sleep(5 * time.Millisecond)
	var req struct {
		Model    string    `json:"model"`
		Messages []Message `json:"messages"`
	}
	if err := json.Unmarshal(body, &req); err != nil || path != "/v1/chat/completions" || req.Model != model {
		return 400, []byte(`{"error":{"message":"bad request"}}`), 0, nil
	}
	f.mu.Lock()
	f.asked = append(f.asked, model)
	f.mu.Unlock()
	last := req.Messages[len(req.Messages)-1].Content
	answer := "I do not know"
	switch {
	case model == "grader":
		answer = `{"score": 0, "reason": "no greeting"}`
		if strings.Contains(last, "Answer:\nHello") {
			answer = `{"score": 1, "reason": "a greeting"}`
		}
	case model == "missing":
		return 404, []byte(`{"error":{"message":"the model \"missing\" does not exist"}}`), 0, nil
	default:
		for k, v := range f.answers {
			if strings.Contains(last, k) {
				answer = v
			}
		}
	}
	out, _ := json.Marshal(map[string]any{"choices": []any{map[string]any{
		"message": map[string]any{"content": answer}, "finish_reason": "stop"}}})
	return 200, out, 0, nil
}

func TestASuiteIsAskedAndScoredByWeight(t *testing.T) {
	s, err := Parse("arith.yaml", []byte(arith))
	if err != nil {
		t.Fatal(err)
	}
	f := &fakeModel{answers: map[string]string{"17 + 25": "<think>carry</think>42", "6 * 7": "41",
		"France": "Paris.", "hello": "Hello there!"}}
	var progress atomic.Int32
	run, err := Evaluate(context.Background(), f.call, s, Options{Model: "qwen", Concurrency: 2,
		Progress: func(done, total int) { progress.Store(int32(done)) }})
	if err != nil {
		t.Fatal(err)
	}
	// add 1, mul 0 at weight 2, capital 1, polite 1: 3 of 5.
	if run.Score != 0.6 || run.Passed || run.Threshold != 0.75 || run.Below != 1 || run.Cases != 4 ||
		run.Suite != "arith" || run.Judge != "grader" || len(run.SuiteHash) != 64 {
		t.Fatalf("%+v", run)
	}
	byID := map[string]store.EvalCase{}
	for _, r := range run.Results {
		byID[r.ID] = r
	}
	if r := byID["add"]; !r.Passed || r.Answer != "42" || r.Kind != "number" {
		t.Errorf("add: %+v", r)
	}
	if r := byID["mul"]; r.Passed || !strings.Contains(r.Detail, "read 41, wanted 42") || r.Weight != 2 {
		t.Errorf("mul: %+v", r)
	}
	if r := byID["polite"]; !r.Passed || !strings.Contains(r.Detail, "the judge: a greeting") {
		t.Errorf("polite: %+v", r)
	}
	if progress.Load() != 4 || f.most.Load() > 2 {
		t.Errorf("progress %d, %d at once with a concurrency of 2", progress.Load(), f.most.Load())
	}

	// A lower bar passes the same answers.
	run, _ = Evaluate(context.Background(), f.call, s, Options{Model: "qwen", MinScore: 0.5})
	if !run.Passed || run.Threshold != 0.5 {
		t.Errorf("at 0.5: %+v", run)
	}
}

func TestAModelThatAnswersNothingFailsWithWhy(t *testing.T) {
	s, _ := Parse("arith.yaml", []byte(arith))
	run, err := Evaluate(context.Background(), (&fakeModel{}).call, s, Options{Model: "missing"})
	if err != nil {
		t.Fatal(err)
	}
	if run.Passed || run.Errors != 4 || !strings.Contains(run.Error, `"missing" does not exist`) {
		t.Fatalf("%+v", run)
	}
}

func TestARubricWithNoJudgeIsRefusedBeforeAnythingIsAsked(t *testing.T) {
	s, _ := Parse("arith.yaml", []byte(strings.Replace(arith, "judge: grader\n", "", 1)))
	f := &fakeModel{}
	if _, err := Evaluate(context.Background(), f.call, s, Options{Model: "qwen"}); err == nil ||
		!strings.Contains(err.Error(), "--judge") {
		t.Fatalf("%v", err)
	}
	if len(f.asked) != 0 {
		t.Errorf("asked %v first", f.asked)
	}
}

func TestASuiteIsReadFromJSONLinesAsFromYAML(t *testing.T) {
	jsonl := `{"name": "fmt", "threshold": 0.5}
{"id": "a", "prompt": "Give JSON", "json_schema": {"type": "object", "required": ["ok"]}}
{"prompt": "Say yes", "regex": "^yes", "ignore_case": true, "weight": 3}
`
	s, err := Parse("format.jsonl", []byte(jsonl))
	if err != nil {
		t.Fatal(err)
	}
	if s.Name != "fmt" || s.Threshold != 0.5 || len(s.Cases) != 2 || s.Cases[1].ID != "case-2" ||
		s.Cases[1].Weight != 3 || s.Cases[0].schema == nil || s.Cases[1].re == nil {
		t.Fatalf("%+v", s)
	}
	// Named by its file when it names nothing, and 0.8 to pass.
	s, err = Parse("dir/plain.jsonl", []byte(`{"prompt": "q", "exact": "a"}`))
	if err != nil || s.Name != "plain" || s.Threshold != DefaultThreshold {
		t.Fatalf("%+v %v", s, err)
	}
}

func TestASuiteThatCannotBeScoredIsRefusedSayingWhy(t *testing.T) {
	for _, tc := range []struct{ suite, why string }{
		{"cases: []", "no cases"},
		{"cases:\n  - prompt: q", "no way to score it"},
		{"cases:\n  - prompt: q\n    exact: a\n    regex: b", "exact and regex"},
		{"cases:\n  - exact: a", "nothing to ask"},
		{"cases:\n  - prompt: q\n    regex: '('", "the regex"},
		{"cases:\n  - id: x\n    prompt: q\n    exact: a\n  - id: x\n    prompt: q\n    exact: a", "two cases are called x"},
		{"threshold: 2\ncases:\n  - prompt: q\n    exact: a", "not 0 to 1"},
		{"cases:\n  - prompt: q\n    exactly: a", "exactly"},
		{"cases:\n  - messages: [{role: robot, content: hi}]\n    exact: a", "role"},
	} {
		if _, err := Parse("s.yaml", []byte(tc.suite)); err == nil || !strings.Contains(err.Error(), tc.why) {
			t.Errorf("%q: %v, wanted it to say %q", tc.suite, err, tc.why)
		}
	}
}

type memStore struct {
	mu   sync.Mutex
	runs []store.EvalRun
}

func (m *memStore) PutEvalRun(_ context.Context, r store.EvalRun) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.runs = append(m.runs, r)
	return nil
}
func (m *memStore) EvalRuns(_ context.Context, model string, _ int) ([]store.EvalRun, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []store.EvalRun
	for i := len(m.runs) - 1; i >= 0; i-- {
		if r := m.runs[i]; model == "" || r.Model == model || r.Target == model {
			r.Results = nil
			out = append(out, r)
		}
	}
	return out, nil
}
func (m *memStore) EvalRunByID(_ context.Context, id string) (store.EvalRun, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, r := range m.runs {
		if r.ID == id {
			return r, nil
		}
	}
	return store.EvalRun{}, store.ErrNotFound
}

func TestTheControlPlaneRunsSuitesByNameAndFilesACanarysRunUnderItsModel(t *testing.T) {
	dir := t.TempDir()
	if err := os.WriteFile(filepath.Join(dir, "arith.yaml"), []byte(arith), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "broken.yaml"), []byte("cases: [}"), 0o600); err != nil {
		t.Fatal(err)
	}
	builtin := fstest.MapFS{
		"arith.yaml":       {Data: []byte("cases:\n  - prompt: shadowed\n    exact: x\n")},
		"json-format.yaml": {Data: []byte("description: JSON\ncases:\n  - prompt: q\n    json_schema: {type: object}\n")},
	}
	f := &fakeModel{answers: map[string]string{"17 + 25": "42", "6 * 7": "42", "France": "Paris", "hello": "Hello"}}
	st := &memStore{}
	svc := NewService(Deps{Call: f.call, Store: st, Dir: dir, Builtin: builtin, Log: slog.New(slog.DiscardHandler),
		Origin: func(m string) (string, string) { return "o/r", "rev-of-" + m },
		CanaryOf: func(n string) string {
			if parent, ok := strings.CutSuffix(n, "-canary"); ok {
				return parent
			}
			return ""
		},
	})

	suites := svc.Suites()
	if len(suites) != 3 || suites[0].Name != "arith" || suites[0].From != filepath.Join(dir, "arith.yaml") ||
		suites[0].Cases != 4 || !suites[0].Judge || suites[1].Error == "" || suites[2].From != "built in" {
		t.Fatalf("%+v", suites)
	}
	if _, err := svc.Load(Request{Suite: "../etc/passwd"}); err == nil {
		t.Error("a path outside the eval directory was read")
	}

	live, err := svc.Start("qwen-canary", Request{Suite: "arith"}, "tester")
	if err != nil {
		t.Fatal(err)
	}
	if live.State != Running || live.Model != "qwen" || live.Target != "qwen-canary" || live.Cases != 4 {
		t.Fatalf("%+v", live)
	}
	deadline := time.Now().Add(5 * time.Second)
	for live.State == Running && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
		live, _ = svc.Get(context.Background(), live.ID)
	}
	if live.State != Done || live.Score != 1 || !live.Passed || live.Revision != "rev-of-qwen-canary" || live.By != "tester" {
		t.Fatalf("%+v", live)
	}
	if kept, err := st.EvalRunByID(context.Background(), live.ID); err != nil || len(kept.Results) != 4 {
		t.Fatalf("kept %+v %v", kept, err)
	}

	// A gate waits for its run, and files it with the job.
	run, err := svc.Run(context.Background(), "qwen", Request{Suite: "json-format"}, "tester", "update-1")
	if err != nil || run.Job != "update-1" || run.Passed || run.Target != "" {
		t.Fatalf("%+v %v", run, err)
	}
	hist, err := svc.History(context.Background(), "qwen", 10)
	if err != nil || len(hist) != 2 || hist[0].ID != run.ID || hist[1].ID != live.ID {
		t.Fatalf("%+v %v", hist, err)
	}
	if _, err := svc.Get(context.Background(), "eval-nope"); !errors.Is(err, store.ErrNotFound) {
		t.Errorf("%v", err)
	}
}