package main

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"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")
	}
}