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

const vertex = createVertex({
  project: process.env.GOOGLE_VERTEX_PROJECT,
  location: process.env.GOOGLE_VERTEX_LOCATION,
});

const customPricingMap: Record<string, ModelPricing> = {
  'gemini-2.5-flash': {
    promptTokens: 0.3 / 1_000_000, // $0.30 per 1M tokens
    completionTokens: 2.5 / 1_000_000, // $2.50 per 1M tokens
    internalReasoningTokens: 2.5 / 1_000_000, // $2.50 per 1M tokens
  },
};

const priceResolver = createObjectPriceResolver(customPricingMap);

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

const wrappedModel = wrapLanguageModel({
  model: vertex('gemini-2.5-flash'),
  middleware: billingMiddleware,
});