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

import { writable, derived } from "svelte/store";
import { invoke } from "@tauri-apps/api/core";
import { notificationManager } from "$lib/notifications/notificationManager";

// Types matching Rust backend
export interface DailyCost {
  date: string;
  input_tokens: number;
  output_tokens: number;
  cost_usd: number;
  messages_sent: number;
  sessions_count: number;
}

export interface CostSummary {
  period_days: number;
  total_input_tokens: number;
  total_output_tokens: number;
  total_cost: number;
  total_messages: number;
  total_sessions: number;
  average_daily_cost: number;
  daily_breakdown: DailyCost[];
}

export type AlertType = "Daily" | "Weekly" | "Monthly";

export interface CostAlert {
  alert_type: AlertType;
  threshold: number;
  current_cost: number;
}

export interface CostAlertThresholds {
  daily: number | null;
  weekly: number | null;
  monthly: number | null;
}

// Store state
interface CostTrackingState {
  todayCost: number;
  weekCost: number;
  monthCost: number;
  summary: CostSummary | null;
  alerts: CostAlert[];
  thresholds: CostAlertThresholds;
  isLoading: boolean;
  lastUpdated: Date | null;
}

const defaultState: CostTrackingState = {
  todayCost: 0,
  weekCost: 0,
  monthCost: 0,
  summary: null,
  alerts: [],
  thresholds: { daily: null, weekly: null, monthly: null },
  isLoading: false,
  lastUpdated: null,
};

function createCostTrackingStore() {
  const { subscribe, set, update } = writable<CostTrackingState>(defaultState);

  return {
    subscribe,

    async refresh() {
      update((s) => ({ ...s, isLoading: true }));

      try {
        const [todayCost, weekCost, monthCost, alerts] = await Promise.all([
          invoke<number>("get_today_cost"),
          invoke<number>("get_week_cost"),
          invoke<number>("get_month_cost"),
          invoke<CostAlert[]>("get_cost_alerts"),
        ]);

        update((s) => ({
          ...s,
          todayCost,
          weekCost,
          monthCost,
          alerts,
          isLoading: false,
          lastUpdated: new Date(),
        }));

        // Trigger notifications for any new alerts
        if (alerts.length > 0) {
          for (const alert of alerts) {
            const message = getAlertMessage(alert);
            notificationManager.notifyCostAlert(message);
          }
        }

        return alerts;
      } catch (error) {
        console.error("Failed to refresh cost tracking:", error);
        update((s) => ({ ...s, isLoading: false }));
        return [];
      }
    },

    async getSummary(days: number): Promise<CostSummary | null> {
      try {
        const summary = await invoke<CostSummary>("get_cost_summary", { days });
        update((s) => ({ ...s, summary }));
        return summary;
      } catch (error) {
        console.error("Failed to get cost summary:", error);
        return null;
      }
    },

    async setAlertThresholds(thresholds: CostAlertThresholds) {
      try {
        await invoke("set_cost_alert_thresholds", {
          daily: thresholds.daily,
          weekly: thresholds.weekly,
          monthly: thresholds.monthly,
        });
        update((s) => ({ ...s, thresholds }));
      } catch (error) {
        console.error("Failed to set alert thresholds:", error);
      }
    },

    async exportCsv(days: number): Promise<string | null> {
      try {
        return await invoke<string>("export_cost_csv", { days });
      } catch (error) {
        console.error("Failed to export CSV:", error);
        return null;
      }
    },

    reset() {
      set(defaultState);
    },
  };
}

export const costTrackingStore = createCostTrackingStore();

// Derived stores for formatted values
export const formattedCosts = derived(costTrackingStore, ($store) => ({
  today: formatCost($store.todayCost),
  week: formatCost($store.weekCost),
  month: formatCost($store.monthCost),
  todayRaw: $store.todayCost,
  weekRaw: $store.weekCost,
  monthRaw: $store.monthCost,
}));

// Helper functions
export function formatCost(cost: number): string {
  if (cost < 0.01) {
    return `$${cost.toFixed(4)}`;
  }
  if (cost < 1) {
    return `$${cost.toFixed(3)}`;
  }
  return `$${cost.toFixed(2)}`;
}

export function formatAlertType(type: AlertType): string {
  switch (type) {
    case "Daily":
      return "Today";
    case "Weekly":
      return "This Week";
    case "Monthly":
      return "This Month";
  }
}

export function getAlertMessage(alert: CostAlert): string {
  const period = formatAlertType(alert.alert_type);
  return `${period}'s spending (${formatCost(alert.current_cost)}) has exceeded your ${formatCost(alert.threshold)} threshold`;
}