package main

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync/atomic"
	"testing"
)

// countingBrain is the LLM Sentry: it counts calls and always answers split with subtasks.
type countingBrain struct{ calls atomic.Int32 }

func (b *countingBrain) Name() string { return "counting" }
func (b *countingBrain) Ask(ctx context.Context, system, prompt string) (string, int, error) {
	b.calls.Add(1)
	return `{"mode":"split","task_type":"writing","subtasks":["a","b"],"reason":"two parts"}`, 10, nil
}

func fakeJev(t *testing.T, status int, mode string, modeConf float64, task string, taskConf float64) *JevClient {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer tk" {
			t.Errorf("auth header = %q", r.Header.Get("Authorization"))
		}
		var body map[string]any
		json.NewDecoder(r.Body).Decode(&body)
		if body["state"] == nil || body["questions"] == nil {
			t.Errorf("request missing state or questions: %v", body)
		}
		w.WriteHeader(status)
		json.NewEncoder(w).Encode(map[string]any{
			"model": "jev-1.13.0",
			"answers": map[string]any{
				"mode":      map[string]any{"type": "choice", "choice": mode, "confidence": modeConf},
				"task_type": map[string]any{"type": "choice", "choice": task, "confidence": taskConf},
			},
			"usage": map[string]int{"input_tokens": 40, "output_tokens": 0},
		})
	}))
	t.Cleanup(srv.Close)
	return &JevClient{url: srv.URL, key: "tk", model: "jev-latest", threshold: 0.7, client: srv.Client()}
}

func triageWith(t *testing.T, j *JevClient) (Triage, int32) {
	t.Helper()
	brain := &countingBrain{}
	f := NewFlock(brain, DefaultGuardrails(), nil)
	f.SetClassifier(j)
	tr, err := f.Triage(context.Background(), nil, "some item")
	if err != nil {
		t.Fatal(err)
	}
	return tr, brain.calls.Load()
}

func TestJevConfidentStandsAlone(t *testing.T) {
	tr, llm := triageWith(t, fakeJev(t, 200, "murder", 0.91, "analysis", 0.88))
	if llm != 0 || tr.Mode != "murder" || tr.TaskType != "analysis" || !strings.Contains(tr.Reason, "Jev: murder 91%") {
		t.Fatalf("got %+v with %d LLM calls; want Jev's call alone", tr, llm)
	}
}

func TestJevUnsureEscalates(t *testing.T) {
	tr, llm := triageWith(t, fakeJev(t, 200, "solo", 0.55, "retrieval", 0.9))
	if llm != 1 || !strings.Contains(tr.Reason, "unsure") {
		t.Fatalf("got %+v with %d LLM calls; want escalation", tr, llm)
	}
}

func TestJevSplitGetsSubtasksFromLLM(t *testing.T) {
	tr, llm := triageWith(t, fakeJev(t, 200, "split", 0.95, "writing", 0.95))
	if llm != 1 || tr.Mode != "split" || len(tr.Subtasks) != 2 || !strings.Contains(tr.Reason, "wrote the subtasks") {
		t.Fatalf("got %+v with %d LLM calls; want LLM-written subtasks", tr, llm)
	}
}

func TestJevDownFallsBackToLLM(t *testing.T) {
	tr, llm := triageWith(t, fakeJev(t, 503, "solo", 0.9, "other", 0.9))
	if llm != 1 || !strings.Contains(tr.Reason, "Jev was unavailable") {
		t.Fatalf("got %+v with %d LLM calls; want fallback", tr, llm)
	}
}

func TestNoJevKeyMeansNoClassifier(t *testing.T) {
	t.Setenv("TYPESAFE_API_KEY", "")
	if NewJev() != nil {
		t.Fatal("expected no classifier without a key")
	}
}

var probeQuestions = map[string]any{
	"kind": map[string]any{
		"type":         "choice",
		"instructions": "What kind of message is this?",
		"criteria":     map[string]string{"ask": "Asks for work", "say": "Ordinary talk", "none": "Neither"},
	},
	"severity": map[string]any{
		"type":         "score",
		"instructions": "How severe is the problem in this message?",
		"criteria":     []string{"No problem", "Annoying", "Blocking"},
	},
	"is_request": map[string]any{
		"type":         "noul",
		"instructions": "Is the sender asking someone to do something?",
	},
}

func TestParseMixedAnswers(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(map[string]any{
			"model": "jev-1.13.0",
			"answers": map[string]any{
				"kind":       map[string]any{"type": "choice", "choice": "ask", "confidence": 0.82},
				"severity":   map[string]any{"type": "score", "score": 1.4, "confidence": 0.6, "legend": map[string]string{"0": "No problem", "1": "Annoying", "2": "Blocking"}, "probabilities": map[string]float64{"0": 0.1, "1": 0.5, "2": 0.4}},
				"is_request": map[string]any{"type": "noul", "noul": 0.95},
			},
			"usage": map[string]int{"input_tokens": 100, "output_tokens": 20},
		})
	}))
	t.Cleanup(srv.Close)
	j := &JevClient{url: srv.URL, key: "tk", model: "jev-latest", client: srv.Client()}
	a, tokens, model, err := j.SystemOne(context.Background(), "state text", probeQuestions)
	if err != nil {
		t.Fatal(err)
	}
	if tokens != 120 || model != "jev-1.13.0" {
		t.Errorf("got tokens=%d model=%s", tokens, model)
	}
	if a["kind"].Choice != "ask" || a["kind"].Confidence == nil || *a["kind"].Confidence != 0.82 {
		t.Errorf("choice answer wrong: %+v", a["kind"])
	}
	if a["severity"].Score == nil || *a["severity"].Score != 1.4 || a["severity"].Legend["2"] != "Blocking" {
		t.Errorf("score answer wrong: %+v", a["severity"])
	}
	if a["is_request"].Noul == nil || *a["is_request"].Noul != 0.95 {
		t.Errorf("noul answer wrong: %+v", a["is_request"])
	}
	if a["is_request"].Confidence != nil {
		t.Errorf("noul must have no confidence, got %v", *a["is_request"].Confidence)
	}
}

func TestModeChoiceHasEscapeOption(t *testing.T) {
	mode := jevQuestions["mode"].(map[string]any)["criteria"].(map[string]string)
	if _, ok := mode["other"]; !ok {
		t.Fatalf("mode has no escape option: %v", mode)
	}
}

func TestClassifySurvivesOtherMode(t *testing.T) {
	tr, _ := triageWith(t, fakeJev(t, 200, "other", 0.95, "analysis", 0.9))
	if tr.Mode != "solo" {
		t.Fatalf("Jev's other mode must normalize to solo, got %q", tr.Mode)
	}
}

func TestProbeScoreAndNoulAgainstLiveAPI(t *testing.T) {
	if os.Getenv("TYPESAFE_API_KEY") == "" {
		t.Skip("no TYPESAFE_API_KEY")
	}
	j := NewJev()
	a, tokens, model, err := j.SystemOne(context.Background(), "@Signals, classify these drops. The export crashes on Safari.", probeQuestions)
	if err != nil {
		t.Fatalf("live probe failed: %v", err)
	}
	dump, _ := json.MarshalIndent(a, "", "  ")
	t.Logf("model=%s tokens=%d\n%s", model, tokens, dump)
	if a["kind"].Choice == "" {
		t.Errorf("choice answer missing: %+v", a["kind"])
	}
	if a["severity"].Score == nil {
		t.Errorf("score answer missing: %+v", a["severity"])
	}
	if a["is_request"].Noul == nil {
		t.Errorf("noul answer missing: %+v", a["is_request"])
	}
}