package main

import (
	"encoding/json"
	"testing"
	"time"
)

func TestExtractUsageHandlesNestedDetailsFromOllama(t *testing.T) {
	tail := []byte(`data: {"id":"x","choices":[],"usage":{"prompt_tokens":35,"prompt_tokens_details":{"cached_tokens":3},"completion_tokens":18,"total_tokens":53}}` + "\n\ndata: [DONE]\n\n")
	p, c, ok := extractUsage(tail)
	if !ok || p != 35 || c != 18 {
		t.Fatalf("got p=%d c=%d ok=%v", p, c, ok)
	}
}

func TestExtractUsageNonStreamBody(t *testing.T) {
	body := []byte(`{"id":"c","choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":7,"completion_tokens":3,"total_tokens":10}}`)
	p, c, ok := extractUsage(body)
	if !ok || p != 7 || c != 3 {
		t.Fatalf("got p=%d c=%d ok=%v", p, c, ok)
	}
}

func TestExtractUsageUsesTheLastOccurrence(t *testing.T) {
	tail := []byte(`data: {"usage":null}` + "\n\n" + `data: {"usage":{"prompt_tokens":1,"completion_tokens":2}}`)
	p, c, ok := extractUsage(tail)
	if !ok || p != 1 || c != 2 {
		t.Fatalf("got p=%d c=%d ok=%v", p, c, ok)
	}
}

func TestExtractUsageMissing(t *testing.T) {
	if _, _, ok := extractUsage([]byte(`data: {"choices":[]}`)); ok {
		t.Fatal("should not find usage")
	}
	if _, _, ok := extractUsage([]byte(`{"usage":null}`)); ok {
		t.Fatal("null usage should not count")
	}
}

func TestHasMatchesExactAndLatest(t *testing.T) {
	set := map[string]bool{"qwen2.5:0.5b": true, "qwen3-0.6b": true, "nvidia/Qwen3.8-27B-NVFP4": true, "imp:latest": true}
	for _, m := range []string{"qwen2.5:0.5b", "qwen3-0.6b", "nvidia/Qwen3.8-27B-NVFP4", "imp", "imp:latest"} {
		if !has(set, m) {
			t.Fatalf("%s should match", m)
		}
	}
	for _, m := range []string{"qwen2.5", "qwen3", "nvidia/other", "imp:v2"} {
		if has(set, m) {
			t.Fatalf("%s should not match", m)
		}
	}
}

func ownerHub(t *testing.T, mode, kind string, retry time.Duration) *Hub {
	t.Helper()
	latest, _ := json.Marshal(map[string]any{"installed": []map[string]any{{"name": "m1", "size": 1, "runtime": "ollama"}}})
	b := &statBuf{latest: latest, at: time.Now(), inst: map[string]bool{"m1": true}, loaded: map[string]bool{}, mode: mode, kind: kind}
	if retry > 0 {
		b.retryAt = time.Now().Add(retry)
	}
	return &Hub{agents: map[string]*conn{"a": {}}, stats: map[string]*statBuf{"a": b}}
}

func TestOwnerWaitReportsKindAndCountdownWithoutProgramNames(t *testing.T) {
	h := ownerHub(t, "owner", "desktop", 4*time.Minute)
	kind, wait := h.ownerWait("m1")
	if kind != "desktop" || wait < 235 || wait > 242 {
		t.Fatalf("kind=%q wait=%d", kind, wait)
	}
	models := h.chatModels()
	if len(models) != 1 || models[0].Available || models[0].Kind != "desktop" || models[0].RetryS < 235 {
		t.Fatalf("%+v", models)
	}
}

func TestOwnerWaitIsUnknownWhenTheMachineCannotSay(t *testing.T) {
	h := ownerHub(t, "owner", "gpu", 0)
	if kind, wait := h.ownerWait("m1"); kind != "gpu" || wait != 0 {
		t.Fatalf("kind=%q wait=%d", kind, wait)
	}
}

func TestAnAvailableMachineHidesTheOwnerInfo(t *testing.T) {
	h := ownerHub(t, "owner", "desktop", time.Minute)
	other := *h.stats["a"]
	other.mode, other.kind, other.retryAt = "available", "", time.Time{}
	h.stats["b"], h.agents["b"] = &other, &conn{}
	models := h.chatModels()
	if len(models) != 1 || !models[0].Available || models[0].Kind != "" || models[0].RetryS != 0 {
		t.Fatalf("%+v", models)
	}
}