hikari-desktop/src/lib/stores/stats.test.ts

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { get } from "svelte/store";
import { listen } from "@tauri-apps/api/event";
import { invoke } from "@tauri-apps/api/core";
import {
  stats,
  formattedStats,
  resetSessionStats,
  contextWarning,
  getContextWarningMessage,
  estimateMessageCost,
  formatTokenCount,
  MODEL_PRICING,
  checkBudget,
  getBudgetStatusMessage,
  getRemainingTokenBudget,
  getRemainingCostBudget,
  initStatsListener,
} from "./stats";
import type { UsageStats, ToolTokenStats, BudgetStatus } from "./stats";

// Helper function to create ToolTokenStats for tests
function toolStats(callCount: number, inputTokens = 0, outputTokens = 0): ToolTokenStats {
  return {
    call_count: callCount,
    estimated_input_tokens: inputTokens,
    estimated_output_tokens: outputTokens,
  };
}

// Mock Tauri APIs
vi.mock("@tauri-apps/api/event", () => ({
  listen: vi.fn(),
}));

vi.mock("@tauri-apps/api/core", () => ({
  invoke: vi.fn(),
}));

vi.mock("./costTracking", () => ({
  costTrackingStore: {
    refresh: vi.fn().mockResolvedValue([]),
  },
}));

describe("stats store", () => {
  beforeEach(() => {
    // Reset stats to default before each test
    stats.set({
      total_input_tokens: 0,
      total_output_tokens: 0,
      total_cost_usd: 0,
      session_input_tokens: 0,
      session_output_tokens: 0,
      session_cost_usd: 0,
      model: null,
      messages_exchanged: 0,
      session_messages_exchanged: 0,
      code_blocks_generated: 0,
      session_code_blocks_generated: 0,
      files_edited: 0,
      session_files_edited: 0,
      files_created: 0,
      session_files_created: 0,
      tools_usage: {},
      session_tools_usage: {},
      session_duration_seconds: 0,
      context_tokens_used: 0,
      context_window_limit: 200000,
      context_utilisation_percent: 0,
      potential_cache_hits: 0,
      potential_cache_savings_tokens: 0,
    });
  });

  describe("stats writable store", () => {
    it("has correct default values", () => {
      const currentStats = get(stats);
      expect(currentStats.total_input_tokens).toBe(0);
      expect(currentStats.total_output_tokens).toBe(0);
      expect(currentStats.total_cost_usd).toBe(0);
      expect(currentStats.model).toBeNull();
    });

    it("can be updated with set", () => {
      const newStats: UsageStats = {
        total_input_tokens: 1000,
        total_output_tokens: 2000,
        total_cost_usd: 0.05,
        session_input_tokens: 500,
        session_output_tokens: 1000,
        session_cost_usd: 0.025,
        model: "claude-sonnet-4",
        messages_exchanged: 10,
        session_messages_exchanged: 5,
        code_blocks_generated: 3,
        session_code_blocks_generated: 2,
        files_edited: 5,
        session_files_edited: 2,
        files_created: 1,
        session_files_created: 1,
        tools_usage: { Read: toolStats(5), Edit: toolStats(3) },
        session_tools_usage: { Read: toolStats(2), Edit: toolStats(1) },
        session_duration_seconds: 300,
        context_tokens_used: 500,
        context_window_limit: 200000,
        context_utilisation_percent: 0.25,
        potential_cache_hits: 0,
        potential_cache_savings_tokens: 0,
      };

      stats.set(newStats);
      const currentStats = get(stats);

      expect(currentStats.total_input_tokens).toBe(1000);
      expect(currentStats.total_output_tokens).toBe(2000);
      expect(currentStats.model).toBe("claude-sonnet-4");
      expect(currentStats.tools_usage.Read?.call_count).toBe(5);
      expect(currentStats.tools_usage.Edit?.call_count).toBe(3);
    });

    it("can be updated with update function", () => {
      stats.update((current) => ({
        ...current,
        total_input_tokens: 500,
        session_messages_exchanged: 3,
      }));

      const currentStats = get(stats);
      expect(currentStats.total_input_tokens).toBe(500);
      expect(currentStats.session_messages_exchanged).toBe(3);
    });
  });

  describe("resetSessionStats", () => {
    it("resets all session fields to zero", () => {
      // First set some values
      stats.set({
        total_input_tokens: 1000,
        total_output_tokens: 2000,
        total_cost_usd: 0.05,
        session_input_tokens: 500,
        session_output_tokens: 1000,
        session_cost_usd: 0.025,
        model: "claude-sonnet-4",
        messages_exchanged: 10,
        session_messages_exchanged: 5,
        code_blocks_generated: 3,
        session_code_blocks_generated: 2,
        files_edited: 5,
        session_files_edited: 2,
        files_created: 1,
        session_files_created: 1,
        tools_usage: { Read: toolStats(5), Edit: toolStats(3) },
        session_tools_usage: { Read: toolStats(2), Edit: toolStats(1) },
        session_duration_seconds: 300,
        context_tokens_used: 500,
        context_window_limit: 200000,
        context_utilisation_percent: 0.25,
        potential_cache_hits: 0,
        potential_cache_savings_tokens: 0,
      });

      // Reset session stats
      resetSessionStats();

      const currentStats = get(stats);

      // Total stats should be preserved
      expect(currentStats.total_input_tokens).toBe(1000);
      expect(currentStats.total_output_tokens).toBe(2000);
      expect(currentStats.total_cost_usd).toBe(0.05);
      expect(currentStats.messages_exchanged).toBe(10);
      expect(currentStats.code_blocks_generated).toBe(3);
      expect(currentStats.files_edited).toBe(5);
      expect(currentStats.files_created).toBe(1);
      expect(currentStats.tools_usage.Read?.call_count).toBe(5);
      expect(currentStats.tools_usage.Edit?.call_count).toBe(3);
      expect(currentStats.model).toBe("claude-sonnet-4");

      // Session stats should be reset
      expect(currentStats.session_input_tokens).toBe(0);
      expect(currentStats.session_output_tokens).toBe(0);
      expect(currentStats.session_cost_usd).toBe(0);
      expect(currentStats.session_messages_exchanged).toBe(0);
      expect(currentStats.session_code_blocks_generated).toBe(0);
      expect(currentStats.session_files_edited).toBe(0);
      expect(currentStats.session_files_created).toBe(0);
      expect(currentStats.session_tools_usage).toEqual({});
      expect(currentStats.session_duration_seconds).toBe(0);
    });
  });

  describe("formattedStats derived store", () => {
    it("formats token numbers with locale string", () => {
      stats.update((current) => ({
        ...current,
        total_input_tokens: 1234567,
        total_output_tokens: 7654321,
        session_input_tokens: 12345,
        session_output_tokens: 54321,
      }));

      const formatted = get(formattedStats);

      expect(formatted.totalTokens).toBe("8,888,888");
      expect(formatted.totalInputTokens).toBe("1,234,567");
      expect(formatted.totalOutputTokens).toBe("7,654,321");
      expect(formatted.sessionTokens).toBe("66,666");
      expect(formatted.sessionInputTokens).toBe("12,345");
      expect(formatted.sessionOutputTokens).toBe("54,321");
    });

    it("formats cost with 4 decimal places", () => {
      stats.update((current) => ({
        ...current,
        total_cost_usd: 1.23456,
        session_cost_usd: 0.00123,
      }));

      const formatted = get(formattedStats);

      expect(formatted.totalCost).toBe("$1.2346");
      expect(formatted.sessionCost).toBe("$0.0012");
    });

    it("formats duration seconds only", () => {
      stats.update((current) => ({
        ...current,
        session_duration_seconds: 45,
      }));

      const formatted = get(formattedStats);
      expect(formatted.sessionDuration).toBe("45s");
    });

    it("formats duration minutes and seconds", () => {
      stats.update((current) => ({
        ...current,
        session_duration_seconds: 125, // 2m 5s
      }));

      const formatted = get(formattedStats);
      expect(formatted.sessionDuration).toBe("2m 5s");
    });

    it("formats duration hours, minutes, and seconds", () => {
      stats.update((current) => ({
        ...current,
        session_duration_seconds: 3725, // 1h 2m 5s
      }));

      const formatted = get(formattedStats);
      expect(formatted.sessionDuration).toBe("1h 2m 5s");
    });

    it("formats duration with zero seconds", () => {
      stats.update((current) => ({
        ...current,
        session_duration_seconds: 3600, // exactly 1h
      }));

      const formatted = get(formattedStats);
      expect(formatted.sessionDuration).toBe("1h 0m 0s");
    });

    it("shows model name when available", () => {
      stats.update((current) => ({
        ...current,
        model: "claude-opus-4-5",
      }));

      const formatted = get(formattedStats);
      expect(formatted.model).toBe("claude-opus-4-5");
    });

    it("shows placeholder when model is null", () => {
      stats.update((current) => ({
        ...current,
        model: null,
      }));

      const formatted = get(formattedStats);
      expect(formatted.model).toBe("No model selected");
    });

    it("formats message counts", () => {
      stats.update((current) => ({
        ...current,
        messages_exchanged: 100,
        session_messages_exchanged: 10,
      }));

      const formatted = get(formattedStats);
      expect(formatted.messagesTotal).toBe("100");
      expect(formatted.messagesSession).toBe("10");
    });

    it("formats code block counts", () => {
      stats.update((current) => ({
        ...current,
        code_blocks_generated: 50,
        session_code_blocks_generated: 5,
      }));

      const formatted = get(formattedStats);
      expect(formatted.codeBlocksTotal).toBe("50");
      expect(formatted.codeBlocksSession).toBe("5");
    });

    it("formats file counts", () => {
      stats.update((current) => ({
        ...current,
        files_edited: 25,
        session_files_edited: 3,
        files_created: 10,
        session_files_created: 2,
      }));

      const formatted = get(formattedStats);
      expect(formatted.filesEditedTotal).toBe("25");
      expect(formatted.filesEditedSession).toBe("3");
      expect(formatted.filesCreatedTotal).toBe("10");
      expect(formatted.filesCreatedSession).toBe("2");
    });

    it("exposes tools usage directly", () => {
      const toolsUsage = { Read: toolStats(10), Edit: toolStats(5), Write: toolStats(3) };
      const sessionToolsUsage = { Read: toolStats(2), Edit: toolStats(1) };

      stats.update((current) => ({
        ...current,
        tools_usage: toolsUsage,
        session_tools_usage: sessionToolsUsage,
      }));

      const formatted = get(formattedStats);
      expect(formatted.toolsUsage).toEqual(toolsUsage);
      expect(formatted.sessionToolsUsage).toEqual(sessionToolsUsage);
    });

    it("handles zero values correctly", () => {
      const formatted = get(formattedStats);

      expect(formatted.totalTokens).toBe("0");
      expect(formatted.totalCost).toBe("$0.0000");
      expect(formatted.sessionDuration).toBe("0s");
      expect(formatted.messagesTotal).toBe("0");
    });

    it("handles large numbers with proper formatting", () => {
      stats.update((current) => ({
        ...current,
        total_input_tokens: 1000000000, // 1 billion
        messages_exchanged: 999999,
      }));

      const formatted = get(formattedStats);
      expect(formatted.totalInputTokens).toBe("1,000,000,000");
      expect(formatted.messagesTotal).toBe("999,999");
    });
  });

  describe("UsageStats interface", () => {
    it("supports all expected fields", () => {
      const fullStats: UsageStats = {
        total_input_tokens: 100,
        total_output_tokens: 200,
        total_cost_usd: 0.01,
        session_input_tokens: 50,
        session_output_tokens: 100,
        session_cost_usd: 0.005,
        model: "test-model",
        messages_exchanged: 5,
        session_messages_exchanged: 2,
        code_blocks_generated: 3,
        session_code_blocks_generated: 1,
        files_edited: 2,
        session_files_edited: 1,
        files_created: 1,
        session_files_created: 0,
        tools_usage: { Read: toolStats(3) },
        session_tools_usage: { Read: toolStats(1) },
        session_duration_seconds: 60,
        context_tokens_used: 50,
        context_window_limit: 200000,
        context_utilisation_percent: 0.025,
        potential_cache_hits: 0,
        potential_cache_savings_tokens: 0,
      };

      stats.set(fullStats);
      const currentStats = get(stats);

      // Verify all fields are present and correct
      expect(currentStats).toEqual(fullStats);
    });
  });

  describe("context window tracking", () => {
    it("tracks context tokens used", () => {
      stats.update((current) => ({
        ...current,
        context_tokens_used: 100000,
        context_window_limit: 200000,
        context_utilisation_percent: 50.0,
      }));

      const currentStats = get(stats);
      expect(currentStats.context_tokens_used).toBe(100000);
      expect(currentStats.context_window_limit).toBe(200000);
      expect(currentStats.context_utilisation_percent).toBe(50.0);
    });

    it("formats context stats correctly", () => {
      stats.update((current) => ({
        ...current,
        context_tokens_used: 150000,
        context_window_limit: 200000,
        context_utilisation_percent: 75.5,
      }));

      const formatted = get(formattedStats);
      expect(formatted.contextUsed).toBe("150,000");
      expect(formatted.contextLimit).toBe("200,000");
      expect(formatted.contextRemaining).toBe("50,000");
      expect(formatted.contextUtilisation).toBe("75.5%");
    });

    it("calculates remaining tokens correctly at limit", () => {
      stats.update((current) => ({
        ...current,
        context_tokens_used: 200000,
        context_window_limit: 200000,
        context_utilisation_percent: 100.0,
      }));

      const formatted = get(formattedStats);
      expect(formatted.contextRemaining).toBe("0");
    });

    it("handles over-limit gracefully", () => {
      stats.update((current) => ({
        ...current,
        context_tokens_used: 250000,
        context_window_limit: 200000,
        context_utilisation_percent: 125.0,
      }));

      const formatted = get(formattedStats);
      expect(formatted.contextRemaining).toBe("0");
    });
  });

  describe("contextWarning derived store", () => {
    it("returns null when under 50%", () => {
      stats.update((current) => ({
        ...current,
        context_utilisation_percent: 40.0,
      }));

      const warning = get(contextWarning);
      expect(warning).toBeNull();
    });

    it("returns moderate when between 50-74%", () => {
      stats.update((current) => ({
        ...current,
        context_utilisation_percent: 60.0,
      }));

      const warning = get(contextWarning);
      expect(warning).toBe("moderate");
    });

    it("returns high when between 75-89%", () => {
      stats.update((current) => ({
        ...current,
        context_utilisation_percent: 80.0,
      }));

      const warning = get(contextWarning);
      expect(warning).toBe("high");
    });

    it("returns critical when 90%+", () => {
      stats.update((current) => ({
        ...current,
        context_utilisation_percent: 95.0,
      }));

      const warning = get(contextWarning);
      expect(warning).toBe("critical");
    });
  });

  describe("getContextWarningMessage", () => {
    it("returns correct message for moderate warning", () => {
      const message = getContextWarningMessage("moderate");
      expect(message).toContain("50%+");
      expect(message).toContain("Consider starting a new conversation");
    });

    it("returns correct message for high warning", () => {
      const message = getContextWarningMessage("high");
      expect(message).toContain("75%+");
      expect(message).toContain("Responses may degrade");
    });

    it("returns correct message for critical warning", () => {
      const message = getContextWarningMessage("critical");
      expect(message).toContain("90%+");
      expect(message).toContain("Start a new conversation");
    });
  });

  describe("formatTokenCount", () => {
    it("formats small numbers directly", () => {
      expect(formatTokenCount(0)).toBe("0");
      expect(formatTokenCount(100)).toBe("100");
      expect(formatTokenCount(999)).toBe("999");
    });

    it("formats thousands with K suffix", () => {
      expect(formatTokenCount(1000)).toBe("1.0K");
      expect(formatTokenCount(1500)).toBe("1.5K");
      expect(formatTokenCount(10000)).toBe("10.0K");
      expect(formatTokenCount(999999)).toBe("1000.0K");
    });

    it("formats millions with M suffix", () => {
      expect(formatTokenCount(1000000)).toBe("1.0M");
      expect(formatTokenCount(1500000)).toBe("1.5M");
      expect(formatTokenCount(10000000)).toBe("10.0M");
    });
  });

  describe("estimateMessageCost", () => {
    it("estimates tokens at ~4 chars per token", () => {
      const result = estimateMessageCost("test", 0, null); // 4 chars = 1 token
      expect(result.messageTokens).toBe(1);
    });

    it("rounds up partial tokens", () => {
      const result = estimateMessageCost("a", 0, null); // 1 char rounds up to 1 token
      expect(result.messageTokens).toBe(1);

      const result2 = estimateMessageCost("abcde", 0, null); // 5 chars = 2 tokens
      expect(result2.messageTokens).toBe(2);
    });

    it("returns 0 tokens for empty string", () => {
      const result = estimateMessageCost("", 0, null);
      expect(result.messageTokens).toBe(0);
      expect(result.estimatedCost).toBe(0);
    });

    it("adds context tokens to total", () => {
      const result = estimateMessageCost("test", 1000, null); // 1 token + 1000 context
      expect(result.messageTokens).toBe(1);
      expect(result.totalInputTokens).toBe(1001);
    });

    it("calculates cost using Sonnet pricing by default", () => {
      // 100 chars = 25 tokens, $3 per million input tokens
      const result = estimateMessageCost("a".repeat(100), 0, null);
      expect(result.messageTokens).toBe(25);
      const expectedCost = (25 / 1_000_000) * 3.0;
      expect(result.estimatedCost).toBeCloseTo(expectedCost, 8);
    });

    it("uses Opus pricing for Opus models", () => {
      const result = estimateMessageCost("a".repeat(100), 0, "claude-opus-4-5-20251101");
      expect(result.messageTokens).toBe(25);
      const expectedCost = (25 / 1_000_000) * 5.0; // Opus 4.5: $5 per million input
      expect(result.estimatedCost).toBeCloseTo(expectedCost, 8);
    });

    it("uses Haiku pricing for Haiku models", () => {
      const result = estimateMessageCost("a".repeat(100), 0, "claude-3-5-haiku-20241022");
      expect(result.messageTokens).toBe(25);
      const expectedCost = (25 / 1_000_000) * 1.0; // Haiku: $1 per million
      expect(result.estimatedCost).toBeCloseTo(expectedCost, 8);
    });

    it("falls back to Sonnet pricing for unknown models", () => {
      const result = estimateMessageCost("a".repeat(100), 0, "unknown-model");
      expect(result.messageTokens).toBe(25);
      const expectedCost = (25 / 1_000_000) * 3.0; // Default Sonnet: $3 per million
      expect(result.estimatedCost).toBeCloseTo(expectedCost, 8);
    });
  });

  describe("MODEL_PRICING", () => {
    it("contains expected Opus pricing", () => {
      // Opus 4.5 has reduced pricing
      expect(MODEL_PRICING["claude-opus-4-5-20251101"]).toEqual({ input: 5.0, output: 25.0 });
      // Previous Opus models have higher pricing
      expect(MODEL_PRICING["claude-opus-4-1-20250805"]).toEqual({ input: 15.0, output: 75.0 });
      expect(MODEL_PRICING["claude-opus-4-20250514"]).toEqual({ input: 15.0, output: 75.0 });
    });

    it("contains expected Sonnet pricing", () => {
      expect(MODEL_PRICING["claude-sonnet-4-5-20250929"]).toEqual({ input: 3.0, output: 15.0 });
      expect(MODEL_PRICING["claude-sonnet-4-20250514"]).toEqual({ input: 3.0, output: 15.0 });
      expect(MODEL_PRICING["claude-3-7-sonnet-20250219"]).toEqual({ input: 3.0, output: 15.0 });
      expect(MODEL_PRICING["claude-3-5-sonnet-20241022"]).toEqual({ input: 3.0, output: 15.0 });
    });

    it("contains expected Haiku pricing", () => {
      expect(MODEL_PRICING["claude-haiku-4-5-20251001"]).toEqual({ input: 1.0, output: 5.0 });
      expect(MODEL_PRICING["claude-3-5-haiku-20241022"]).toEqual({ input: 1.0, output: 5.0 });
      expect(MODEL_PRICING["claude-3-haiku-20240307"]).toEqual({ input: 0.25, output: 1.25 });
    });
  });

  describe("checkBudget", () => {
    const baseStats: UsageStats = {
      total_input_tokens: 0,
      total_output_tokens: 0,
      total_cost_usd: 0,
      session_input_tokens: 500,
      session_output_tokens: 500,
      session_cost_usd: 0.5,
      model: null,
      messages_exchanged: 0,
      session_messages_exchanged: 0,
      code_blocks_generated: 0,
      session_code_blocks_generated: 0,
      files_edited: 0,
      session_files_edited: 0,
      files_created: 0,
      session_files_created: 0,
      tools_usage: {},
      session_tools_usage: {},
      session_duration_seconds: 0,
      context_tokens_used: 0,
      context_window_limit: 200000,
      context_utilisation_percent: 0,
      potential_cache_hits: 0,
      potential_cache_savings_tokens: 0,
    };

    it("returns ok when budget is disabled", () => {
      const result = checkBudget(baseStats, false, 100, 1.0, 0.8);
      expect(result).toEqual({ type: "ok" });
    });

    it("returns ok when under all budgets", () => {
      const result = checkBudget(baseStats, true, 10000, 10.0, 0.8);
      expect(result).toEqual({ type: "ok" });
    });

    it("returns exceeded when token budget is reached", () => {
      const result = checkBudget(baseStats, true, 1000, null, 0.8);
      expect(result).toEqual({ type: "exceeded", budget_type: "token" });
    });

    it("returns warning when token usage is above threshold", () => {
      // session tokens = 1000, budget = 1100, threshold = 0.8 → 1000/1100 ≈ 0.909 > 0.8
      const result = checkBudget(baseStats, true, 1100, null, 0.8);
      expect(result.type).toBe("warning");
      if (result.type === "warning") {
        expect(result.budget_type).toBe("token");
        expect(result.percent_used).toBeGreaterThan(80);
      }
    });

    it("returns exceeded when cost budget is reached", () => {
      const result = checkBudget(baseStats, true, null, 0.5, 0.8);
      expect(result).toEqual({ type: "exceeded", budget_type: "cost" });
    });

    it("returns warning when cost usage is above threshold", () => {
      // session cost = 0.5, budget = 0.55, threshold = 0.8 → 0.5/0.55 ≈ 0.909 > 0.8
      const result = checkBudget(baseStats, true, null, 0.55, 0.8);
      expect(result.type).toBe("warning");
      if (result.type === "warning") {
        expect(result.budget_type).toBe("cost");
      }
    });

    it("checks token budget before cost budget", () => {
      // Both budgets exceeded, but token should be reported first
      const result = checkBudget(baseStats, true, 500, 0.1, 0.8);
      expect(result).toEqual({ type: "exceeded", budget_type: "token" });
    });

    it("returns ok when no budgets are set", () => {
      const result = checkBudget(baseStats, true, null, null, 0.8);
      expect(result).toEqual({ type: "ok" });
    });
  });

  describe("getBudgetStatusMessage", () => {
    it("returns null for ok status", () => {
      const status: BudgetStatus = { type: "ok" };
      expect(getBudgetStatusMessage(status)).toBeNull();
    });

    it("returns exceeded message for token budget", () => {
      const status: BudgetStatus = { type: "exceeded", budget_type: "token" };
      const message = getBudgetStatusMessage(status);
      expect(message).toContain("token");
      expect(message).toContain("exceeded");
    });

    it("returns exceeded message for cost budget", () => {
      const status: BudgetStatus = { type: "exceeded", budget_type: "cost" };
      const message = getBudgetStatusMessage(status);
      expect(message).toContain("cost");
      expect(message).toContain("exceeded");
    });

    it("returns warning message with percentage for token budget", () => {
      const status: BudgetStatus = { type: "warning", budget_type: "token", percent_used: 85.5 };
      const message = getBudgetStatusMessage(status);
      expect(message).toContain("token");
      expect(message).toContain("86%");
    });

    it("returns warning message with percentage for cost budget", () => {
      const status: BudgetStatus = { type: "warning", budget_type: "cost", percent_used: 90 };
      const message = getBudgetStatusMessage(status);
      expect(message).toContain("cost");
      expect(message).toContain("90%");
    });
  });

  describe("getRemainingTokenBudget", () => {
    const baseStats: UsageStats = {
      total_input_tokens: 0,
      total_output_tokens: 0,
      total_cost_usd: 0,
      session_input_tokens: 300,
      session_output_tokens: 200,
      session_cost_usd: 0,
      model: null,
      messages_exchanged: 0,
      session_messages_exchanged: 0,
      code_blocks_generated: 0,
      session_code_blocks_generated: 0,
      files_edited: 0,
      session_files_edited: 0,
      files_created: 0,
      session_files_created: 0,
      tools_usage: {},
      session_tools_usage: {},
      session_duration_seconds: 0,
      context_tokens_used: 0,
      context_window_limit: 200000,
      context_utilisation_percent: 0,
      potential_cache_hits: 0,
      potential_cache_savings_tokens: 0,
    };

    it("returns null when no token budget is set", () => {
      expect(getRemainingTokenBudget(baseStats, null)).toBeNull();
    });

    it("returns the remaining tokens when under budget", () => {
      // session tokens = 500, budget = 1000 → remaining = 500
      expect(getRemainingTokenBudget(baseStats, 1000)).toBe(500);
    });

    it("returns 0 when at or over budget", () => {
      expect(getRemainingTokenBudget(baseStats, 500)).toBe(0);
      expect(getRemainingTokenBudget(baseStats, 400)).toBe(0);
    });
  });

  describe("getRemainingCostBudget", () => {
    const baseStats: UsageStats = {
      total_input_tokens: 0,
      total_output_tokens: 0,
      total_cost_usd: 0,
      session_input_tokens: 0,
      session_output_tokens: 0,
      session_cost_usd: 0.3,
      model: null,
      messages_exchanged: 0,
      session_messages_exchanged: 0,
      code_blocks_generated: 0,
      session_code_blocks_generated: 0,
      files_edited: 0,
      session_files_edited: 0,
      files_created: 0,
      session_files_created: 0,
      tools_usage: {},
      session_tools_usage: {},
      session_duration_seconds: 0,
      context_tokens_used: 0,
      context_window_limit: 200000,
      context_utilisation_percent: 0,
      potential_cache_hits: 0,
      potential_cache_savings_tokens: 0,
    };

    it("returns null when no cost budget is set", () => {
      expect(getRemainingCostBudget(baseStats, null)).toBeNull();
    });

    it("returns the remaining cost when under budget", () => {
      // session cost = 0.3, budget = 1.0 → remaining = 0.7
      expect(getRemainingCostBudget(baseStats, 1.0)).toBeCloseTo(0.7, 5);
    });

    it("returns 0 when at or over budget", () => {
      expect(getRemainingCostBudget(baseStats, 0.3)).toBe(0);
      expect(getRemainingCostBudget(baseStats, 0.2)).toBe(0);
    });
  });

  describe("initStatsListener", () => {
    const emptyStats: UsageStats = {
      total_input_tokens: 0,
      total_output_tokens: 0,
      total_cost_usd: 0,
      session_input_tokens: 0,
      session_output_tokens: 0,
      session_cost_usd: 0,
      model: null,
      messages_exchanged: 0,
      session_messages_exchanged: 0,
      code_blocks_generated: 0,
      session_code_blocks_generated: 0,
      files_edited: 0,
      session_files_edited: 0,
      files_created: 0,
      session_files_created: 0,
      tools_usage: {},
      session_tools_usage: {},
      session_duration_seconds: 0,
      context_tokens_used: 0,
      context_window_limit: 200000,
      context_utilisation_percent: 0,
      potential_cache_hits: 0,
      potential_cache_savings_tokens: 0,
    };

    afterEach(() => {
      vi.resetAllMocks();
    });

    it("registers a listener for the claude:stats event", async () => {
      vi.mocked(listen).mockResolvedValue(vi.fn());
      vi.mocked(invoke).mockResolvedValue(emptyStats);

      await initStatsListener();

      expect(listen).toHaveBeenCalledWith("claude:stats", expect.any(Function));
    });

    it("updates the stats store when the listener callback fires", async () => {
      let capturedCallback: ((event: unknown) => void) | undefined;
      vi.mocked(listen).mockImplementation(async (_event, callback) => {
        capturedCallback = callback as (event: unknown) => void;
        return () => {};
      });
      vi.mocked(invoke).mockResolvedValue(emptyStats);

      await initStatsListener();

      const newStats = { ...emptyStats, total_input_tokens: 9999 };
      capturedCallback!({ payload: { stats: newStats } });

      expect(get(stats).total_input_tokens).toBe(9999);
    });

    it("loads persisted stats from the backend on initialisation", async () => {
      const persistedStats = { ...emptyStats, total_input_tokens: 5000 };
      vi.mocked(listen).mockResolvedValue(vi.fn());
      vi.mocked(invoke).mockResolvedValue(persistedStats);

      await initStatsListener();

      expect(invoke).toHaveBeenCalledWith("get_persisted_stats");
      expect(get(stats).total_input_tokens).toBe(5000);
    });

    it("handles a failed invoke gracefully and logs the error", async () => {
      const error = new Error("Failed to get stats");
      vi.mocked(listen).mockResolvedValue(vi.fn());
      vi.mocked(invoke).mockRejectedValue(error);

      const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {});

      await initStatsListener();

      expect(consoleErrorSpy).toHaveBeenCalledWith("Failed to load initial stats:", error);
      consoleErrorSpy.mockRestore();
    });
  });
});