import { createGroq } from '@ai-sdk/groq';
import { wrapLanguageModel } from 'ai';
import { createGroqMiddleware } from '@ai-billing/groq';
import {
  consoleDestination,
  createObjectPriceResolver,
  type ModelPricing,
} from '@ai-billing/core';

const groq = createGroq({ apiKey: process.env.GROQ_API_KEY });

const customPricingMap: Record<string, ModelPricing> = {
  'openai/gpt-oss-120b': {
    promptTokens: 0.15 / 1_000_000,
    completionTokens: 0.6 / 1_000_000,
    inputCacheReadTokens: 0.075 / 1_000_000,
    inputCacheWriteTokens: 0,
  },
};

const priceResolver = createObjectPriceResolver(customPricingMap);

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

const wrappedModel = wrapLanguageModel({
  model: groq('openai/gpt-oss-120b'),
  middleware: billingMiddleware,
});