import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock';
import { wrapLanguageModel } from 'ai';
import { createAmazonBedrockV4Middleware } from '@ai-billing/amazon-bedrock';
import {
  consoleDestination,
  createObjectPriceResolver,
  type ModelPricing,
} from '@ai-billing/core';

const amazonBedrock = createAmazonBedrock({
  region: process.env.AWS_REGION,
});

const customPricingMap: Record<string, ModelPricing> = {
  'us.amazon.nova-lite-v1:0': {
    promptTokens: 0.06 / 1_000_000,
    completionTokens: 0.24 / 1_000_000,
  },
};

const priceResolver = createObjectPriceResolver(customPricingMap);

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

const wrappedModel = wrapLanguageModel({
  model: amazonBedrock('us.amazon.nova-lite-v1:0'),
  middleware: billingMiddleware,
});