import z from "zod";
import { defineTool } from "../../core/define-tool";
import {
  requireAuth,
  requireOrganization,
  getUserId,
} from "../../core/studio-context";
import { HOSTED_PROVIDER_IDS } from "../../ai-providers/provider-ids";
import { getProviders } from "../../ai-providers/registry";
import { mintGatewayJwt } from "../../auth/jwt";

/** Guards against a malformed upstream gateway response (missing/NaN/Infinity
 *  balance) silently reaching billing UI as an unusable number. */
export function isValidBalanceCents(balanceCents: number): boolean {
  return Number.isFinite(balanceCents);
}

export const AI_PROVIDER_CREDITS = defineTool({
  name: "AI_PROVIDER_CREDITS",
  description:
    "Get the current credit balance for a provider (providers that support it, e.g. Deco AI Gateway)",
  inputSchema: z.object({
    providerId: z.enum(HOSTED_PROVIDER_IDS),
  }),
  outputSchema: z.object({
    balanceCents: z
      .number()
      .describe("Remaining balance in cents (e.g. 1000 = $10.00)"),
  }),
  handler: async (input, ctx) => {
    requireAuth(ctx);
    const org = requireOrganization(ctx);
    await ctx.access.check();

    const userId = getUserId(ctx);
    if (!userId) throw new Error("Unable to determine user ID");

    const adapter = getProviders()[input.providerId];
    if (!adapter) {
      throw new Error(`Unknown provider: ${input.providerId}`);
    }
    if (!adapter.getCreditsBalance) {
      throw new Error(
        `Provider ${input.providerId} does not expose a credits balance`,
      );
    }

    const studioJwt = await mintGatewayJwt(userId);

    const result = await adapter.getCreditsBalance(studioJwt, org.id);
    if (!isValidBalanceCents(result.balanceCents)) {
      console.error(
        `AI provider "${input.providerId}" returned a non-finite credits balance for org ${org.id}: ${result.balanceCents}`,
      );
      throw new Error(
        `Provider ${input.providerId} returned an invalid credits balance`,
      );
    }

    return result;
  },
});
