/** biome-ignore-all lint/style/noMagicNumbers: we're defining constants */
import type { LanguageModelUsage } from "ai";
import { saveUsage } from "@/utils/redis/usage";
import { OPENROUTER_MODEL_PRICING } from "@/utils/llms/pricing.generated";
import {
  STATIC_MODEL_PRICING,
  type ModelPricing,
} from "@/utils/llms/supported-model-pricing";
import {
  getOpenRouterProviderPrefix,
  stripOnlineModelSuffix,
} from "@/utils/llms/model-id";
import { publishAiCall } from "@inboxzero/tinybird-ai-analytics";
import { createScopedLogger } from "@/utils/logger";

const logger = createScopedLogger("usage");

export type AiUsageEvent = {
  cachedInputTokens: number;
  estimatedCost: number;
  inputTokens: number;
  label: string;
  model: string;
  outputTokens: number;
  platformCost: number;
  provider: string;
  providerCostSource?: string;
  providerReportedCost?: number;
  providerUpstreamInferenceCost?: number;
  reasoningTokens: number;
  totalTokens: number;
};

type AiUsageListener = (event: AiUsageEvent) => void;

const aiUsageListeners = new Set<AiUsageListener>();

export function subscribeToAiUsage(listener: AiUsageListener): () => void {
  aiUsageListeners.add(listener);
  return () => {
    aiUsageListeners.delete(listener);
  };
}

export async function saveAiUsage({
  userId,
  email,
  emailAccountId,
  provider,
  model,
  usage,
  label,
  hasUserApiKey,
  providerReportedCost,
  providerUpstreamInferenceCost,
  providerCostSource,
  stepCount,
  toolCallCount,
}: {
  userId?: string;
  email: string;
  emailAccountId: string;
  provider: string;
  model: string;
  usage: LanguageModelUsage;
  label: string;
  hasUserApiKey?: boolean;
  providerReportedCost?: number;
  providerUpstreamInferenceCost?: number;
  providerCostSource?: string;
  stepCount?: number;
  toolCallCount?: number;
}) {
  const estimatedCost = calculateUsageCost({ provider, model, usage });
  const isUserApiKey = !!hasUserApiKey;
  const platformCost = isUserApiKey ? 0 : estimatedCost;
  const inputTokens = usage.inputTokens ?? 0;
  const outputTokens = usage.outputTokens ?? 0;
  const cachedInputTokens = usage.cachedInputTokens ?? 0;
  const reasoningTokens = usage.reasoningTokens ?? 0;
  const totalTokens = usage.totalTokens ?? 0;

  notifyAiUsageListeners({
    cachedInputTokens,
    estimatedCost,
    inputTokens,
    label,
    model,
    outputTokens,
    platformCost,
    provider,
    providerCostSource,
    providerReportedCost,
    providerUpstreamInferenceCost,
    reasoningTokens,
    totalTokens,
  });

  try {
    return Promise.all([
      publishAiCall({
        userId: userId ?? email,
        emailAccountId,
        provider,
        model,
        totalTokens,
        completionTokens: outputTokens,
        promptTokens: inputTokens,
        cachedInputTokens,
        reasoningTokens,
        cost: platformCost,
        estimatedCost,
        providerReportedCost,
        providerUpstreamInferenceCost,
        providerCostSource,
        isUserApiKey: toTinybirdBoolean(isUserApiKey),
        timestamp: Date.now(),
        label,
        stepCount,
        toolCallCount,
      }),
      saveUsage({ userId, emailAccountId, cost: platformCost, usage }),
    ]);
  } catch (error) {
    logger.error("Failed to save usage", { error });
  }
}

export function calculateUsageCost(options: {
  provider: string;
  model: string;
  usage: LanguageModelUsage;
}): number {
  const { provider, model, usage } = options;
  const pricing = getModelPricing({ provider, model });
  if (!pricing) return 0;

  const rawCachedInputTokens = usage.cachedInputTokens ?? 0;
  const normalizedCachedInputTokens = Math.max(0, rawCachedInputTokens);
  const inputTokens = Math.max(
    0,
    usage.inputTokens ?? normalizedCachedInputTokens,
  );
  const cachedInputTokens = Math.min(inputTokens, normalizedCachedInputTokens);
  const uncachedInputTokens = Math.max(0, inputTokens - cachedInputTokens);
  const outputTokens = Math.max(0, usage.outputTokens ?? 0);
  const reasoningTokens = Math.max(0, usage.reasoningTokens ?? 0);
  const cachedInputTokenPrice = pricing.cachedInput ?? pricing.input;

  return (
    uncachedInputTokens * pricing.input +
    cachedInputTokens * cachedInputTokenPrice +
    (outputTokens + reasoningTokens) * pricing.output
  );
}

function getModelPricing(options: {
  provider: string;
  model: string;
}): ModelPricing | undefined {
  const { provider, model } = options;
  const providerId = provider.toLowerCase();

  for (const candidate of buildModelLookupCandidates({
    model,
    provider: providerId,
  })) {
    if (providerId === "openrouter") {
      const openRouterPricing = OPENROUTER_MODEL_PRICING[candidate];
      if (openRouterPricing) return openRouterPricing;
    }

    const fallbackPricing = STATIC_MODEL_PRICING[candidate];
    if (fallbackPricing) return fallbackPricing;

    if (providerId !== "openrouter") {
      const openRouterPricing = OPENROUTER_MODEL_PRICING[candidate];
      if (openRouterPricing) return openRouterPricing;
    }
  }

  return;
}

function buildModelLookupCandidates({
  provider,
  model,
}: {
  provider: string;
  model: string;
}): string[] {
  const noOnlineSuffix = stripOnlineModelSuffix(model);

  const candidates = [model, noOnlineSuffix];
  const unprefixed = noOnlineSuffix.includes("/")
    ? noOnlineSuffix.split("/").at(-1)
    : null;

  if (unprefixed) {
    candidates.push(unprefixed);
  } else {
    const providerPrefix = getOpenRouterProviderPrefix(provider);
    if (providerPrefix) {
      candidates.push(`${providerPrefix}/${noOnlineSuffix}`);
    }
  }

  return [...new Set(candidates)];
}

function notifyAiUsageListeners(event: AiUsageEvent): void {
  for (const listener of aiUsageListeners) {
    try {
      listener(event);
    } catch (error) {
      logger.error("AI usage listener failed", { error });
    }
  }
}

function toTinybirdBoolean(value: boolean): 0 | 1 {
  return value ? 1 : 0;
}