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

import { describe, it, expect, beforeEach, vi } from "vitest";
import { get } from "svelte/store";
import {
  stats,
  formattedStats,
  resetSessionStats,
  contextWarning,
  getContextWarningMessage,
  estimateMessageCost,
  formatTokenCount,
  MODEL_PRICING,
} from "./stats";
import type { UsageStats, ToolTokenStats } 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(),
}));

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 });
    });
  });
});