package bot

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
)

func rowsByKey(rows []UsageRow) map[string]int64 {
	out := map[string]int64{}
	for _, r := range rows {
		out[r.Day+"|"+r.Model+"|"+r.Dimension] = r.Quantity
	}
	return out
}

// The accumulator's contract: calls merge into (day, model, dimension)
// cells, drain empties, and restore folds a failed flush back in losslessly.
func TestUsageMeterAccumulateDrainRestore(t *testing.T) {
	m := newUsageMeter()
	m.addLLM("claude-sonnet-5", anthUsage{InputTokens: 100, OutputTokens: 20, CacheReadInputTokens: 300})
	m.addLLM("claude-sonnet-5", anthUsage{InputTokens: 50, OutputTokens: 5, CacheCreationInputTokens: 7})
	m.addLLM("claude-haiku-4-5", anthUsage{InputTokens: 9})
	m.inc("charter_runs", 1)

	day := meterDay()
	rows := m.drain()
	got := rowsByKey(rows)
	for key, want := range map[string]int64{
		day + "|claude-sonnet-5|llm_requests":           2,
		day + "|claude-sonnet-5|llm_input_tokens":       150,
		day + "|claude-sonnet-5|llm_output_tokens":      25,
		day + "|claude-sonnet-5|llm_cache_read_tokens":  300,
		day + "|claude-sonnet-5|llm_cache_write_tokens": 7,
		day + "|claude-haiku-4-5|llm_requests":          1,
		day + "|claude-haiku-4-5|llm_input_tokens":      9,
		day + "||charter_runs":                          1,
	} {
		if got[key] != want {
			t.Errorf("%s = %d, want %d", key, got[key], want)
		}
	}
	if m.drain() != nil {
		t.Fatal("second drain should be empty")
	}

	// A failed flush restores; new counts merge on top.
	m.restore(rows)
	m.addLLM("claude-sonnet-5", anthUsage{InputTokens: 1})
	again := rowsByKey(m.drain())
	if again[day+"|claude-sonnet-5|llm_input_tokens"] != 151 {
		t.Fatalf("restore+add did not merge: %v", again)
	}
}

// A zero-usage block (a failed parse, an empty response) records only the
// request count — never zero-quantity rows.
func TestUsageMeterSkipsZeroes(t *testing.T) {
	m := newUsageMeter()
	m.addLLM("m", anthUsage{})
	rows := m.drain()
	if len(rows) != 1 || rows[0].Dimension != "llm_requests" || rows[0].Quantity != 1 {
		t.Fatalf("want just llm_requests=1, got %+v", rows)
	}
}

// The non-streamed path: the response's usage block reaches both the returned
// response and the record hook (the meter's tap).
func TestAnthropicUsageParsed(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(map[string]any{
			"content":     []map[string]any{{"type": "text", "text": "hi"}},
			"stop_reason": "end_turn",
			"usage": map[string]any{
				"input_tokens": 120, "output_tokens": 34,
				"cache_creation_input_tokens": 8, "cache_read_input_tokens": 4000,
			},
		})
	}))
	defer srv.Close()

	var gotModel string
	var got anthUsage
	c := &anthropicClient{key: "k", base: srv.URL, http: srv.Client(),
		record: func(model string, u anthUsage) { gotModel, got = model, u }}
	resp, err := c.create(anthRequest{Model: "claude-sonnet-5", MaxTokens: 10})
	if err != nil {
		t.Fatalf("create: %v", err)
	}
	want := anthUsage{InputTokens: 120, OutputTokens: 34, CacheCreationInputTokens: 8, CacheReadInputTokens: 4000}
	if resp.Usage != want {
		t.Fatalf("response usage = %+v, want %+v", resp.Usage, want)
	}
	if gotModel != "claude-sonnet-5" || got != want {
		t.Fatalf("record hook got (%q, %+v)", gotModel, got)
	}
}

// The streamed path: input + cache counters ride message_start, the final
// (cumulative) output count rides message_delta.
func TestAnthropicUsageParsedStream(t *testing.T) {
	stream := "event: message_start\n" +
		`data: {"type":"message_start","message":{"usage":{"input_tokens":200,"output_tokens":1,"cache_read_input_tokens":900}}}` + "\n\n" +
		"event: content_block_start\n" +
		`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n" +
		"event: content_block_delta\n" +
		`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}` + "\n\n" +
		"event: message_delta\n" +
		`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":57}}` + "\n\n" +
		"event: message_stop\n" +
		`data: {"type":"message_stop"}` + "\n\n"
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte(stream))
	}))
	defer srv.Close()

	var got anthUsage
	c := &anthropicClient{key: "k", base: srv.URL, http: srv.Client(),
		record: func(_ string, u anthUsage) { got = u }}
	resp, err := c.createStream(anthRequest{Model: "m", MaxTokens: 10})
	if err != nil {
		t.Fatalf("createStream: %v", err)
	}
	want := anthUsage{InputTokens: 200, OutputTokens: 57, CacheReadInputTokens: 900}
	if resp.Usage != want {
		t.Fatalf("stream usage = %+v, want %+v", resp.Usage, want)
	}
	if got != want {
		t.Fatalf("record hook got %+v", got)
	}
}

// The outbox contract (the ack-loss fix): a failed flush keeps the report
// VERBATIM — same id, same rows — so the ledger can dedupe the replay, and
// usage counted during the outage rides the NEXT report, never the pending one.
func TestFlushUsageOutboxRetriesVerbatim(t *testing.T) {
	type report struct {
		ReportID string     `json:"report_id"`
		Rows     []UsageRow `json:"rows"`
	}
	var got []report
	fail := true
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var rep report
		json.NewDecoder(r.Body).Decode(&rep)
		got = append(got, rep)
		if fail {
			http.Error(w, "cloud blip", http.StatusBadGateway)
			return
		}
		json.NewEncoder(w).Encode(map[string]any{"ok": true, "budget": map[string]any{"state": "ok"}})
	}))
	defer srv.Close()

	r := &Runtime{Base: srv.URL, Session: &Session{Token: "tok", Name: "Eleven"}, Logf: t.Logf}
	r.client = srv.Client()
	r.meter = newUsageMeter()

	r.meter.inc("charter_runs", 1)
	r.flushUsage() // fails; report stays pending
	r.meter.inc("charter_runs", 5)
	r.flushUsage() // still failing; the SAME report is resent, new counts stay in the meter
	fail = false
	r.flushUsage() // succeeds: pending report lands
	r.flushUsage() // then the accumulated counts land as a NEW report

	if len(got) != 4 {
		t.Fatalf("attempts = %d, want 4", len(got))
	}
	if got[0].ReportID == "" || got[0].ReportID != got[1].ReportID || got[1].ReportID != got[2].ReportID {
		t.Fatalf("retries must reuse the report id: %q %q %q", got[0].ReportID, got[1].ReportID, got[2].ReportID)
	}
	if got[3].ReportID == got[2].ReportID {
		t.Fatal("post-success flush must mint a fresh report id")
	}
	if len(got[2].Rows) != 1 || got[2].Rows[0].Quantity != 1 {
		t.Fatalf("pending report mutated across retries: %+v", got[2].Rows)
	}
	if len(got[3].Rows) != 1 || got[3].Rows[0].Quantity != 5 {
		t.Fatalf("outage-time counts must ride the next report: %+v", got[3].Rows)
	}
	if r.pendingUsage != nil {
		t.Fatal("nothing should be pending after success")
	}
	if r.budgetState.Load() != "ok" {
		t.Fatalf("budget state not stored: %v", r.budgetState.Load())
	}
}

// Phase 3's sidecar half: an "exhausted" budget state gates the model paths
// (one honest notice, then silence) and recovery re-arms the notice.
func TestBudgetExhaustionGate(t *testing.T) {
	r := &Runtime{Session: &Session{Name: "Eleven"}}
	if r.budgetExhausted() {
		t.Fatal("never-flushed must read as within budget")
	}
	r.budgetState.Store("exhausted")
	if !r.budgetExhausted() {
		t.Fatal("exhausted state not read")
	}
	if !r.noteBudgetExhausted() {
		t.Fatal("first notice must fire")
	}
	if r.noteBudgetExhausted() {
		t.Fatal("the notice is once-per-exhaustion")
	}
	// A task run refuses without burning tokens (ai is set, budget is not).
	r.ai = &anthropicClient{}
	if _, _, err := r.execTaskOnce("p", &TaskSpec{Prompt: "x"}, false); err == nil {
		t.Fatal("exhausted budget must refuse task runs")
	}
	// Recovery (simulated flush storing a non-exhausted state) re-arms.
	r.budgetState.Store("ok")
	r.budgetNotified.Store(false)
	if r.budgetExhausted() {
		t.Fatal("ok state must clear the gate")
	}
	if !r.noteBudgetExhausted() {
		t.Fatal("notice must re-arm after recovery")
	}
}