import {
  clearCache,
  readCache,
  writeCache,
  type EvenStorage,
} from "./live-cache";
import {
  LEGACY_AI_PRICING_VERSION,
  emptyAiUsageCharge,
  mergeAiUsageCharge,
  type AiUsageCharge,
} from "./ai-pricing";

export type AiUsage = {
  readonly textInputTokens: number;
  readonly cachedTextInputTokens: number;
  readonly audioInputTokens: number;
  readonly cachedAudioInputTokens: number;
  readonly textOutputTokens: number;
  readonly transcriptionAudioInputTokens: number;
  readonly transcriptionTextOutputTokens: number;
  readonly searchTextInputTokens: number;
  readonly cachedSearchTextInputTokens: number;
  readonly searchTextOutputTokens: number;
  readonly webSearchCalls: number;
};

export type AiDailyUsage = {
  readonly date: string;
  readonly usage: AiUsage;
  readonly charge: AiUsageCharge;
};

export const EMPTY_AI_USAGE: AiUsage = {
  textInputTokens: 0,
  cachedTextInputTokens: 0,
  audioInputTokens: 0,
  cachedAudioInputTokens: 0,
  textOutputTokens: 0,
  transcriptionAudioInputTokens: 0,
  transcriptionTextOutputTokens: 0,
  searchTextInputTokens: 0,
  cachedSearchTextInputTokens: 0,
  searchTextOutputTokens: 0,
  webSearchCalls: 0,
};

const PRICE_PER_MILLION = {
  textInputTokens: 4,
  cachedTextInputTokens: 0.4,
  audioInputTokens: 32,
  cachedAudioInputTokens: 0.4,
  textOutputTokens: 16,
  transcriptionAudioInputTokens: 1.25,
  transcriptionTextOutputTokens: 5,
  searchTextInputTokens: 0.25,
  cachedSearchTextInputTokens: 0,
  searchTextOutputTokens: 2,
  webSearchCalls: 10_000,
} as const satisfies Record<keyof AiUsage, number>;

function safeTokenCount(value: number): number {
  return Number.isFinite(value) ? Math.max(0, Math.floor(value)) : 0;
}

export function addAiUsage(left: AiUsage, right: AiUsage): AiUsage {
  return Object.fromEntries(
    (Object.keys(EMPTY_AI_USAGE) as Array<keyof AiUsage>).map((key) => [
      key,
      safeTokenCount(left[key]) + safeTokenCount(right[key]),
    ]),
  ) as unknown as AiUsage;
}

export function estimateAiUsageUsd(usage: AiUsage): number {
  return (Object.keys(PRICE_PER_MILLION) as Array<keyof AiUsage>)
    .reduce(
      (total, key) => total
        + safeTokenCount(usage[key]) * PRICE_PER_MILLION[key] / 1_000_000,
      0,
    );
}

function localDateKey(value: Date): string {
  return [
    value.getFullYear(),
    String(value.getMonth() + 1).padStart(2, "0"),
    String(value.getDate()).padStart(2, "0"),
  ].join("-");
}

export function addDailyAiUsage(
  ledger: readonly AiDailyUsage[],
  at: Date,
  usage: AiUsage,
  charge: AiUsageCharge = legacyCharge(usage),
): readonly AiDailyUsage[] {
  const date = localDateKey(at);
  const previousEntry = ledger.find((entry) => entry.date === date);
  return [
    ...ledger.filter((entry) => entry.date !== date),
    {
      date,
      usage: addAiUsage(previousEntry?.usage ?? EMPTY_AI_USAGE, usage),
      charge: mergeAiUsageCharge(
        previousEntry?.charge ?? emptyAiUsageCharge(),
        charge,
      ),
    },
  ].sort((left, right) => left.date.localeCompare(right.date)).slice(-400);
}

function legacyCharge(usage: AiUsage): AiUsageCharge {
  return {
    estimatedNanoUsd: Math.round(estimateAiUsageUsd(usage) * 1_000_000_000),
    unpricedEvents: 0,
    pricingVersions: [LEGACY_AI_PRICING_VERSION],
  };
}

function usageInRange(
  ledger: readonly AiDailyUsage[],
  start: Date,
  end: Date,
): AiUsage {
  const startKey = localDateKey(start);
  const endKey = localDateKey(end);
  return ledger.reduce(
    (total, entry) => entry.date >= startKey && entry.date <= endKey
      ? addAiUsage(total, entry.usage)
      : total,
    EMPTY_AI_USAGE,
  );
}

export function usageForCurrentWeek(
  ledger: readonly AiDailyUsage[],
  now = new Date(),
): AiUsage {
  const start = new Date(now.getFullYear(), now.getMonth(), now.getDate());
  const day = start.getDay();
  start.setDate(start.getDate() - (day === 0 ? 6 : day - 1));
  return usageInRange(ledger, start, now);
}

export function usageForCurrentMonth(
  ledger: readonly AiDailyUsage[],
  now = new Date(),
): AiUsage {
  const start = new Date(now.getFullYear(), now.getMonth(), 1);
  return usageInRange(ledger, start, now);
}

function chargeInRange(
  ledger: readonly AiDailyUsage[],
  start: Date,
  end: Date,
): AiUsageCharge {
  const startKey = localDateKey(start);
  const endKey = localDateKey(end);
  return ledger.reduce(
    (total, entry) => entry.date >= startKey && entry.date <= endKey
      ? mergeAiUsageCharge(total, entry.charge)
      : total,
    emptyAiUsageCharge(),
  );
}

export function costSummaryForCurrentPeriod(
  ledger: readonly AiDailyUsage[],
  now = new Date(),
): { weekUsd: number; monthUsd: number; hasUnpricedUsage: boolean } {
  const weekStart = new Date(now.getFullYear(), now.getMonth(), now.getDate());
  const day = weekStart.getDay();
  weekStart.setDate(weekStart.getDate() - (day === 0 ? 6 : day - 1));
  const monthStart = new Date(now.getFullYear(), now.getMonth(), 1);
  const week = chargeInRange(ledger, weekStart, now);
  const month = chargeInRange(ledger, monthStart, now);
  return {
    weekUsd: week.estimatedNanoUsd / 1_000_000_000,
    monthUsd: month.estimatedNanoUsd / 1_000_000_000,
    hasUnpricedUsage: week.unpricedEvents > 0 || month.unpricedEvents > 0,
  };
}

const LEGACY_USAGE_KEYS = Object.keys(EMPTY_AI_USAGE).filter((key) => (
  !key.toLowerCase().includes("search")
));

function isUsage(value: unknown): value is AiUsage {
  if (typeof value !== "object" || value === null) return false;
  const item = value as Record<string, unknown>;
  return LEGACY_USAGE_KEYS.every(
    (key) => typeof item[key] === "number"
      && Number.isFinite(item[key])
      && (item[key] as number) >= 0,
  );
}

function isUsageLedger(value: unknown): value is readonly AiDailyUsage[] {
  return Array.isArray(value)
    && value.length <= 400
    && value.every((entry) => (
      typeof entry === "object"
      && entry !== null
      && typeof entry.date === "string"
      && /^\d{4}-\d{2}-\d{2}$/.test(entry.date)
      && isUsage(entry.usage)
      && (
        entry.charge === undefined
        || isCharge(entry.charge)
      )
    ));
}

function isCharge(value: unknown): value is AiUsageCharge {
  if (typeof value !== "object" || value === null) return false;
  const item = value as Record<string, unknown>;
  return typeof item.estimatedNanoUsd === "number"
    && Number.isFinite(item.estimatedNanoUsd)
    && item.estimatedNanoUsd >= 0
    && typeof item.unpricedEvents === "number"
    && Number.isFinite(item.unpricedEvents)
    && item.unpricedEvents >= 0
    && Array.isArray(item.pricingVersions)
    && item.pricingVersions.every((version) => typeof version === "string");
}

export async function resolveAiUsageLedger(
  storage: EvenStorage,
): Promise<readonly AiDailyUsage[]> {
  const ledger = await readCache(storage, "ai-usage", isUsageLedger) ?? [];
  return ledger.map((entry) => ({
    ...entry,
    usage: addAiUsage(EMPTY_AI_USAGE, entry.usage),
    charge: entry.charge ?? legacyCharge(addAiUsage(EMPTY_AI_USAGE, entry.usage)),
  }));
}

export function writeAiUsageLedger(
  storage: EvenStorage,
  ledger: readonly AiDailyUsage[],
): Promise<boolean> {
  return writeCache(storage, "ai-usage", ledger.slice(-400));
}

export function clearAiUsageLedger(storage: EvenStorage): Promise<boolean> {
  return clearCache(storage, "ai-usage");
}