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

import { writable, derived } from "svelte/store";
import { listen } from "@tauri-apps/api/event";
import { invoke } from "@tauri-apps/api/core";
import { costTrackingStore } from "./costTracking";
import { configStore } from "./config";

export type ContextWarning = "moderate" | "high" | "critical";
export type BudgetType = "token" | "cost";

// Model pricing (per million tokens) - keep in sync with stats.rs
// Source: https://platform.claude.com/docs/en/about-claude/models/overview
export const MODEL_PRICING: Record<string, { input: number; output: number }> = {
  // Current generation (Claude 4.6)
  "claude-opus-4-6": { input: 5.0, output: 25.0 },
  "claude-sonnet-4-6": { input: 3.0, output: 15.0 },
  // Previous generation (Claude 4.5)
  "claude-opus-4-5-20251101": { input: 5.0, output: 25.0 },
  "claude-sonnet-4-5-20250929": { input: 3.0, output: 15.0 },
  "claude-haiku-4-5-20251001": { input: 1.0, output: 5.0 },
  // Previous generation (Claude 4.x)
  "claude-opus-4-1-20250805": { input: 15.0, output: 75.0 },
  "claude-opus-4-20250514": { input: 15.0, output: 75.0 },
  "claude-sonnet-4-20250514": { input: 3.0, output: 15.0 },
  // Legacy (Claude 3.x)
  "claude-3-7-sonnet-20250219": { input: 3.0, output: 15.0 },
  "claude-3-5-sonnet-20241022": { input: 3.0, output: 15.0 },
  "claude-3-5-sonnet-20240620": { input: 3.0, output: 15.0 },
  "claude-3-5-haiku-20241022": { input: 1.0, output: 5.0 },
  "claude-3-opus-20240229": { input: 15.0, output: 75.0 },
  "claude-3-sonnet-20240229": { input: 3.0, output: 15.0 },
  "claude-3-haiku-20240307": { input: 0.25, output: 1.25 },
};

const DEFAULT_PRICING = { input: 3.0, output: 15.0 }; // Default to Sonnet

export interface CostEstimate {
  messageTokens: number;
  totalInputTokens: number;
  estimatedCost: number;
}

// Estimate cost for a message before sending
export function estimateMessageCost(
  messageText: string,
  contextTokensUsed: number,
  model: string | null
): CostEstimate {
  // Estimate tokens using ~4 chars per token heuristic
  const messageTokens = Math.ceil(messageText.length / 4);
  const totalInputTokens = contextTokensUsed + messageTokens;

  const pricing = model ? (MODEL_PRICING[model] ?? DEFAULT_PRICING) : DEFAULT_PRICING;
  const estimatedCost = (totalInputTokens / 1_000_000) * pricing.input;

  return { messageTokens, totalInputTokens, estimatedCost };
}
export type BudgetStatus =
  | { type: "ok" }
  | { type: "warning"; budget_type: BudgetType; percent_used: number }
  | { type: "exceeded"; budget_type: BudgetType };

// Per-tool token usage statistics
export interface ToolTokenStats {
  call_count: number;
  estimated_input_tokens: number;
  estimated_output_tokens: number;
}

export interface UsageStats {
  total_input_tokens: number;
  total_output_tokens: number;
  total_cost_usd: number;
  session_input_tokens: number;
  session_output_tokens: number;
  session_cost_usd: number;
  model: string | null;

  // New fields
  messages_exchanged: number;
  session_messages_exchanged: number;
  code_blocks_generated: number;
  session_code_blocks_generated: number;
  files_edited: number;
  session_files_edited: number;
  files_created: number;
  session_files_created: number;
  tools_usage: Record<string, ToolTokenStats>;
  session_tools_usage: Record<string, ToolTokenStats>;
  session_duration_seconds: number;

  // Context window tracking
  context_tokens_used: number;
  context_window_limit: number;
  context_utilisation_percent: number;

  // Cache analytics (tracks potential savings from repeated tool calls)
  potential_cache_hits: number;
  potential_cache_savings_tokens: number;
}

// Main stats store
export const stats = writable<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,
});

// Format token count with K/M suffix
export function formatTokenCount(tokens: number): string {
  if (tokens >= 1000000) {
    return `${(tokens / 1000000).toFixed(1)}M`;
  }
  if (tokens >= 1000) {
    return `${(tokens / 1000).toFixed(1)}K`;
  }
  return tokens.toString();
}

// Derived store for formatted display values
export const formattedStats = derived([stats, configStore.config], ([$stats, $config]) => {
  const formatNumber = (num: number) => num.toLocaleString();
  const formatCost = (cost: number) => `$${cost.toFixed(4)}`;
  const formatDuration = (seconds: number) => {
    const hours = Math.floor(seconds / 3600);
    const minutes = Math.floor((seconds % 3600) / 60);
    const secs = seconds % 60;

    if (hours > 0) {
      return `${hours}h ${minutes}m ${secs}s`;
    } else if (minutes > 0) {
      return `${minutes}m ${secs}s`;
    } else {
      return `${secs}s`;
    }
  };

  // Format tool stats with token info
  const formatToolStats = (toolStats: Record<string, ToolTokenStats>) => {
    return Object.entries(toolStats).map(([name, stats]) => ({
      name,
      callCount: stats.call_count,
      totalTokens: stats.estimated_input_tokens + stats.estimated_output_tokens,
      formattedTokens: formatTokenCount(
        stats.estimated_input_tokens + stats.estimated_output_tokens
      ),
      inputTokens: stats.estimated_input_tokens,
      outputTokens: stats.estimated_output_tokens,
    }));
  };

  // Use the model from stats if available, otherwise fall back to the configured model
  const currentModel = $stats.model ?? $config.model ?? "No model selected";

  return {
    totalTokens: formatNumber($stats.total_input_tokens + $stats.total_output_tokens),
    totalInputTokens: formatNumber($stats.total_input_tokens),
    totalOutputTokens: formatNumber($stats.total_output_tokens),
    totalCost: formatCost($stats.total_cost_usd),
    sessionTokens: formatNumber($stats.session_input_tokens + $stats.session_output_tokens),
    sessionInputTokens: formatNumber($stats.session_input_tokens),
    sessionOutputTokens: formatNumber($stats.session_output_tokens),
    sessionCost: formatCost($stats.session_cost_usd),
    model: currentModel,

    // New formatted fields
    messagesTotal: formatNumber($stats.messages_exchanged),
    messagesSession: formatNumber($stats.session_messages_exchanged),
    codeBlocksTotal: formatNumber($stats.code_blocks_generated),
    codeBlocksSession: formatNumber($stats.session_code_blocks_generated),
    filesEditedTotal: formatNumber($stats.files_edited),
    filesEditedSession: formatNumber($stats.session_files_edited),
    filesCreatedTotal: formatNumber($stats.files_created),
    filesCreatedSession: formatNumber($stats.session_files_created),
    sessionDuration: formatDuration($stats.session_duration_seconds),
    toolsUsage: $stats.tools_usage,
    sessionToolsUsage: $stats.session_tools_usage,
    // Formatted tool stats with token info
    sessionToolsFormatted: formatToolStats($stats.session_tools_usage),
    toolsFormatted: formatToolStats($stats.tools_usage),

    // Context window tracking
    contextUsed: formatNumber($stats.context_tokens_used),
    contextLimit: formatNumber($stats.context_window_limit),
    contextRemaining: formatNumber(
      Math.max(0, $stats.context_window_limit - $stats.context_tokens_used)
    ),
    contextUtilisation: `${$stats.context_utilisation_percent.toFixed(1)}%`,
  };
});

// Derived store for context warning state
export const contextWarning = derived(stats, ($stats): ContextWarning | null => {
  if ($stats.context_utilisation_percent >= 90) {
    return "critical";
  } else if ($stats.context_utilisation_percent >= 75) {
    return "high";
  } else if ($stats.context_utilisation_percent >= 50) {
    return "moderate";
  }
  return null;
});

// Get warning message for context utilisation
export function getContextWarningMessage(warning: ContextWarning): string {
  switch (warning) {
    case "moderate":
      return "Context window is 50%+ full. Consider starting a new conversation for better performance.";
    case "high":
      return "Context window is 75%+ full. Responses may degrade. Consider summarising or starting fresh.";
    case "critical":
      return "Context window is nearly full (90%+)! Start a new conversation to avoid errors.";
  }
}

// Budget checking functions
export function checkBudget(
  stats: UsageStats,
  budgetEnabled: boolean,
  tokenBudget: number | null,
  costBudget: number | null,
  warningThreshold: number
): BudgetStatus {
  if (!budgetEnabled) {
    return { type: "ok" };
  }

  const sessionTokens = stats.session_input_tokens + stats.session_output_tokens;

  // Check token budget
  if (tokenBudget !== null) {
    if (sessionTokens >= tokenBudget) {
      return { type: "exceeded", budget_type: "token" };
    }
    const percentUsed = sessionTokens / tokenBudget;
    if (percentUsed >= warningThreshold) {
      return { type: "warning", budget_type: "token", percent_used: percentUsed * 100 };
    }
  }

  // Check cost budget
  if (costBudget !== null) {
    if (stats.session_cost_usd >= costBudget) {
      return { type: "exceeded", budget_type: "cost" };
    }
    const percentUsed = stats.session_cost_usd / costBudget;
    if (percentUsed >= warningThreshold) {
      return { type: "warning", budget_type: "cost", percent_used: percentUsed * 100 };
    }
  }

  return { type: "ok" };
}

// Get budget status message
export function getBudgetStatusMessage(status: BudgetStatus): string | null {
  if (status.type === "ok") {
    return null;
  }

  const budgetTypeLabel = status.budget_type === "token" ? "token" : "cost";

  if (status.type === "exceeded") {
    return `Session ${budgetTypeLabel} budget exceeded! Consider starting a new session.`;
  }

  return `Approaching ${budgetTypeLabel} budget limit (${status.percent_used.toFixed(0)}% used).`;
}

// Get remaining budget values
export function getRemainingTokenBudget(
  stats: UsageStats,
  tokenBudget: number | null
): number | null {
  if (tokenBudget === null) return null;
  const used = stats.session_input_tokens + stats.session_output_tokens;
  return Math.max(0, tokenBudget - used);
}

export function getRemainingCostBudget(
  stats: UsageStats,
  costBudget: number | null
): number | null {
  if (costBudget === null) return null;
  return Math.max(0, costBudget - stats.session_cost_usd);
}

// Note: Cost calculation is now done in the Rust backend

// Initialize stats listener
export async function initStatsListener() {
  // Listen for stats updates from the backend
  await listen("claude:stats", (event) => {
    const payload = event.payload as { stats: UsageStats };
    const { stats: newStats } = payload;

    // The backend already tracks all totals - just set the stats directly
    stats.set(newStats);

    // Refresh cost tracking to check for alerts (debounced - won't spam)
    costTrackingStore.refresh();
  });

  // Load initial persisted stats from backend (no bridge required)
  try {
    const initialStats = await invoke<UsageStats>("get_persisted_stats");
    stats.set(initialStats);
    console.log("Loaded persisted stats:", initialStats);
  } catch (error) {
    console.error("Failed to load initial stats:", error);
  }
}

// Reset session stats (call when starting new session)
export function resetSessionStats() {
  stats.update((current) => ({
    ...current,
    session_input_tokens: 0,
    session_output_tokens: 0,
    session_cost_usd: 0,
    session_messages_exchanged: 0,
    session_code_blocks_generated: 0,
    session_files_edited: 0,
    session_files_created: 0,
    session_tools_usage: {},
    session_duration_seconds: 0,
  }));
}