import { createGoogleGenerativeAI } from '@ai-sdk/google';
import { wrapLanguageModel } from 'ai';
import { createGoogleV4Middleware } from '@ai-billing/google';
import {
  consoleDestination,
  createObjectPriceResolver,
  type ModelPricing,
} from '@ai-billing/core';

const google = createGoogleGenerativeAI({ apiKey: process.env.GOOGLE_AI_STUDIO_KEY });

const customPricingMap: Record<string, ModelPricing> = {
 'models/gemini-3.1-flash-lite-preview': {
   promptTokens: 0.00000025, // $0.25 per 1M tokens
   completionTokens: 0.0000015, // $1.50 per 1M tokens
   inputCacheReadTokens: 0.000000025, // $0.025 per 1M tokens
   internalReasoningTokens: 0.0000015, // $1.50 per 1M tokens
 },
};

const priceResolver = createObjectPriceResolver(customPricingMap);

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

const wrappedModel = wrapLanguageModel({
  model: google('models/gemini-3.1-flash-lite-preview'),
  middleware: billingMiddleware,
});