pieced-portal/src/app/api/usage/route.ts

import { NextRequest, NextResponse } from "next/server";
import { getSessionUser } from "@/lib/session";
import { listTenants } from "@/lib/k8s";
import { getTeamInfo, getTeamSpendLogsV2 } from "@/lib/litellm";
import { safeError } from "@/lib/errors";

/**
 * GET /api/usage
 *
 * Customers: teamId is resolved server-side from the tenant matching the
 *            user's orgId. No client-supplied teamId accepted.
 * Platform admins: may pass ?teamId=... to inspect any tenant's usage.
 */
export async function GET(req: NextRequest) {
  const user = await getSessionUser();
  if (!user)
    return NextResponse.json({ error: "Unauthorized" }, { status: 401 });

  let teamId: string | null = null;

  if (user.isPlatform) {
    // Admins may pass a specific teamId to inspect any tenant
    teamId = req.nextUrl.searchParams.get("teamId") ?? null;
  }

  // For customers (or admins without explicit teamId): resolve from their tenant
  if (!teamId) {
    const tenants = await listTenants();
    const orgTenant = tenants.find(
      (t) => t.metadata.labels?.["pieced.ch/zitadel-org-id"] === user.orgId
    );

    if (!orgTenant?.status?.litellmTeamId) {
      return NextResponse.json(
        { error: "No active tenant found for your organization" },
        { status: 404 }
      );
    }
    teamId = orgTenant.status.litellmTeamId;
  }

  // Month param: YYYY-MM, defaults to current month
  const now = new Date();
  const monthParam =
    req.nextUrl.searchParams.get("month") ||
    `${now.getFullYear()}-${String(now.getMonth() + 1).padStart(2, "0")}`;

  const [year, month] = monthParam.split("-").map(Number);
  const startDate = new Date(year, month - 1, 1);
  const endDate = new Date(year, month, 0); // last day of month

  const startStr = startDate.toISOString().split("T")[0];
  const endStr = endDate.toISOString().split("T")[0];

  try {
    const teamInfo = await getTeamInfo(teamId);

    // Fetch all pages
    const allRequests: any[] = [];
    let page = 1;
    while (true) {
      const result = await getTeamSpendLogsV2(
        teamId,
        startStr,
        endStr,
        page,
        100
      );
      allRequests.push(...(result.data || []));
      if (page >= (result.total_pages || 1)) break;
      page++;
    }

    // Aggregate by day
    const byDay: Record<
      string,
      { inputTokens: number; outputTokens: number; spend: number }
    > = {};
    for (const r of allRequests) {
      const day = (r.startTime || r.endTime || "").slice(0, 10);
      if (!day) continue;
      if (!byDay[day])
        byDay[day] = { inputTokens: 0, outputTokens: 0, spend: 0 };
      byDay[day].inputTokens += r.prompt_tokens || 0;
      byDay[day].outputTokens += r.completion_tokens || 0;
      byDay[day].spend += r.spend || 0;
    }

    const dailyUsage = Object.entries(byDay)
      .sort(([a], [b]) => a.localeCompare(b))
      .map(([date, d]) => ({ date, ...d }));

    const totalInput = allRequests.reduce(
      (s, r) => s + (r.prompt_tokens || 0),
      0
    );
    const totalOutput = allRequests.reduce(
      (s, r) => s + (r.completion_tokens || 0),
      0
    );
    const totalSpend = allRequests.reduce((s, r) => s + (r.spend || 0), 0);

    return NextResponse.json({
      teamId,
      month: monthParam,
      currentPeriod: {
        inputTokens: totalInput,
        outputTokens: totalOutput,
        totalSpend,
        requestCount: allRequests.length,
      },
      budget: {
        maxBudget: teamInfo?.team_info?.max_budget ?? null,
        spend: teamInfo?.team_info?.spend ?? 0,
        remaining: teamInfo?.team_info?.max_budget
          ? teamInfo.team_info.max_budget - (teamInfo.team_info.spend ?? 0)
          : null,
      },
      rateLimits: {
        rpm: teamInfo?.team_info?.rpm_limit ?? null,
        tpm: teamInfo?.team_info?.tpm_limit ?? null,
      },
      dailyUsage,
    });
  } catch (e: any) {
    console.error("Usage fetch error:", e.message);
    return NextResponse.json(
      { error: safeError(e, "Failed to fetch usage") },
      { status: 500 }
    );
  }
}