import { createPerplexity } from '@ai-sdk/perplexity';
import { wrapLanguageModel } from 'ai';
import { createPerplexityV3Middleware } from '@ai-billing/perplexity';
import {
  consoleDestination,
  createObjectPriceResolver,
  type ModelPricing,
} from '@ai-billing/core';

const perplexity = createPerplexity({ apiKey: process.env.PERPLEXITY_API_KEY });

const customPricingMap: Record<string, ModelPricing> = {
  'sonar-pro': {
    promptTokens: 3.0 / 1_000_000,
    completionTokens: 15.0 / 1_000_000,
    request: 6.0 / 1_000, // flat per-request search fee (low search_context_size rate)
  },
};

const priceResolver = createObjectPriceResolver(customPricingMap);

const billingMiddleware = createPerplexityV3Middleware({
  destinations: [consoleDestination()],
  priceResolver,
});

const wrappedModel = wrapLanguageModel({
  model: perplexity('sonar-pro'),
  middleware: billingMiddleware,
});