import { createOpenAICompatible } from '@ai-sdk/openai-compatible';
import { wrapLanguageModel } from 'ai';
import { createOpenAICompatibleV4Middleware } from '@ai-billing/openai-compatible';
import {
  consoleDestination,
  createObjectPriceResolver,
  type ModelPricing,
} from '@ai-billing/core';
import type { LanguageModelV4 } from '@ai-sdk/provider';

const xai = createOpenAICompatible({
  name: 'xai', // must match `providerId` below
  baseURL: 'https://api.x.ai/v1',
  apiKey: process.env.XAI_API_KEY,
  includeUsage: true,
});

const customPricingMap: Record<string, ModelPricing> = {
  'grok-4-1-fast-reasoning': {
    promptTokens: 0.2 / 1_000_000,
    completionTokens: 0.5 / 1_000_000,
    internalReasoningTokens: 0,
    inputCacheReadTokens: 0.15 / 1_000_000,
  },
};

const priceResolver = createObjectPriceResolver(customPricingMap);

const billingMiddleware = createOpenAICompatibleV4Middleware({
  destinations: [consoleDestination()],
  priceResolver,
  providerId: 'xai',
});

const wrappedModel = wrapLanguageModel({
  model: xai('grok-4-1-fast-reasoning') as unknown as LanguageModelV4,
  middleware: billingMiddleware,
});