import { UsagePilot } from "usagepilot";
import { PostgresStorage } from "usagepilot/storage";

// Define your dimension structure
type Dimensions = {
  user_id: string;
  team_id: string;
  organization_id: string;
};

// Create the storage adapter
const storage = new PostgresStorage({
  connectionString: process.env.DATABASE_URL,
  // Optional configuration
  batchSize: 100,
  flushInterval: 30000, // 30 seconds
  maxRetries: 3,
  retryDelay: 1000,
});

// Initialize UsagePilot
const usage = new UsagePilot<Dimensions>({
  storage,
});

import { streamText } from "ai";
import { createOpenAI } from "@ai-sdk/openai";

// Create a fetch wrapper that tracks usage
const trackedFetch = usage.createFetch((req) => ({
  user_id: req.headers.get("x-user-id"),
  team_id: req.headers.get("x-team-id"),
  organization_id: req.headers.get("x-organization-id"),
}));

// Create OpenAI instance with tracked fetch
const openai = createOpenAI({
  apiKey: process.env.OPENAI_API_KEY,
  fetch: trackedFetch,
});

// Use it with Vercel AI SDK
async function generateResponse(prompt: string, userId: string) {
  const result = await streamText({
    model: openai("gpt-4"),
    prompt,
    headers: {
      "x-user-id": userId,
      "x-team-id": "team_456",
      "x-organization-id": "org_123",
    },
  });

  // Usage is automatically tracked and stored!
  return result;
}

// Get usage for a specific organization this month
const startOfMonth = new Date();
startOfMonth.setDate(1);
startOfMonth.setHours(0, 0, 0, 0);

const result = await usage.query({
  dimensions: {
    organization_id: "org_123",
  },
  timeRange: {
    start: startOfMonth,
    end: new Date(),
  },
});

// Calculate total tokens used
const totalTokens = result.aggregates.reduce(
  (sum, agg) => sum + agg.metrics.totalTokens,
  0,
);

console.log(`Organization used ${totalTokens} tokens this month`);

// Get daily usage for billing
const dailyUsage = await usage.query({
  dimensions: {
    organization_id: "org_123",
  },
  timeRange: {
    start: billingPeriodStart,
    end: billingPeriodEnd,
  },
  granularity: "day", // Aggregate by day
});

// Calculate total tokens from daily aggregates
const totalTokens = dailyUsage.aggregates.reduce(
  (sum, agg) => sum + agg.metrics.totalTokens,
  0,
);

// Send to your billing provider
await stripe.usageRecords.create({
  subscription_item: "si_123",
  quantity: totalTokens,
  timestamp: Math.floor(Date.now() / 1000),
});