import { createAzure } from '@ai-sdk/azure';
import { wrapLanguageModel } from 'ai';
import { createAzureV4Middleware } from '@ai-billing/azure';
import {
  consoleDestination,
  createObjectPriceResolver,
  type ModelPricing,
} from '@ai-billing/core';

const azure = createAzure({
  apiKey: process.env.AZURE_API_KEY,
  baseURL: `${process.env.AZURE_URL?.replace(//+$/, '')}/openai/v1`,
  apiVersion: 'preview',
});

const deployment = process.env.AZURE_DEPLOYMENT ?? 'DeepSeek-V4-Pro';

const customPricingMap: Record<string, ModelPricing> = {
  [deployment]: {
    promptTokens: 2.0 / 1_000_000,
    completionTokens: 6.0 / 1_000_000,
  },
};

const priceResolver = createObjectPriceResolver(customPricingMap);

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

const wrappedModel = wrapLanguageModel({
  model: azure(deployment),
  middleware: billingMiddleware,
});