import { NextRequest, NextResponse } from "next/server";
import { z } from "zod";
import { getSessionUser, canMutate } from "@/lib/session";
import { getTenant } from "@/lib/k8s";
import { canUserSeeTenant } from "@/lib/visibility";
import { findKeyByAlias, updateKeyBudget } from "@/lib/litellm";
import { safeError } from "@/lib/errors";

/**
 * Update the per-tenant budget — operates on the LiteLLM virtual
 * key, NOT on the team.
 *
 * Why per-key
 * -----------
 * Each tenant in an org has its own virtual key
 * (`key_alias = tenant.metadata.name`); the team that owns those
 * keys is org-scoped and shared across all the org's tenants. A
 * budget on the team would cap the whole org; a budget on the key
 * caps just this one tenant. Customers landing on the tenant detail
 * page reasonably expect "edit budget" to mean "the budget of THIS
 * tenant" — so we put it on the key.
 *
 * The team-level (org-wide) budget is a separate control that lives
 * in /settings (not yet implemented) — the two coexist: LiteLLM
 * applies whichever cap is hit first.
 *
 * Schema:
 *   - maxBudget: number > 0 (set a cap), or null (remove the cap).
 *   - budgetDuration: one of "30d", "1mo", "1y", or null (lifetime).
 *
 * Authorization: owners and platform admins.
 */

const patchSchema = z.object({
  // > 0 because LiteLLM rejects 0 and a zero cap would lock the key
  // out instantly. Upper bound 1M as a typo guard.
  maxBudget: z.number().positive().max(1_000_000).nullable(),
  budgetDuration: z.enum(["30d", "1mo", "1y"]).nullable(),
});

export async function PATCH(
  req: NextRequest,
  { params }: { params: Promise<{ name: string }> }
) {
  const user = await getSessionUser();
  if (!user) {
    return NextResponse.json({ error: "Unauthorized" }, { status: 401 });
  }
  if (!canMutate(user)) {
    return NextResponse.json({ error: "Forbidden" }, { status: 403 });
  }

  const { name } = await params;
  const tenant = await getTenant(name);
  if (!tenant) {
    return NextResponse.json({ error: "Not found" }, { status: 404 });
  }
  if (!(await canUserSeeTenant(user, tenant))) {
    // Don't leak existence — same 404 a non-visible tenant gets.
    return NextResponse.json({ error: "Not found" }, { status: 404 });
  }

  const teamId = tenant.status?.litellmTeamId;
  if (!teamId) {
    return NextResponse.json(
      {
        error:
          "Tenant has no LiteLLM team yet. Please wait until provisioning completes.",
      },
      { status: 409 }
    );
  }

  const body = await req.json().catch(() => null);
  const parsed = patchSchema.safeParse(body);
  if (!parsed.success) {
    return NextResponse.json(
      { error: "Invalid input", details: parsed.error.flatten() },
      { status: 400 }
    );
  }

  // Defensive: removing the cap should null out the duration too —
  // a reset cadence on an unlimited budget is meaningless and would
  // confuse LiteLLM's bookkeeping.
  const maxBudget = parsed.data.maxBudget;
  const budgetDuration =
    maxBudget === null ? null : parsed.data.budgetDuration;

  // Look up the key by alias (= tenant name). The token returned is
  // what /key/update wants in the `key` field.
  let keyInfo;
  try {
    keyInfo = await findKeyByAlias(teamId, name);
  } catch (e: any) {
    console.error("Failed to look up tenant key:", e);
    return NextResponse.json(
      { error: safeError(e, "Failed to look up tenant key") },
      { status: 500 }
    );
  }
  if (!keyInfo) {
    return NextResponse.json(
      {
        error:
          "Tenant has no virtual key yet. Please wait until provisioning completes.",
      },
      { status: 409 }
    );
  }

  try {
    await updateKeyBudget(keyInfo.token, { maxBudget, budgetDuration });
    return NextResponse.json({
      message: maxBudget === null ? "Budget removed." : "Budget updated.",
      maxBudget,
      budgetDuration,
    });
  } catch (e: any) {
    console.error("Failed to update key budget:", e);
    return NextResponse.json(
      { error: safeError(e, "Failed to update budget") },
      { status: 500 }
    );
  }
}