import type { Tool } from "ai";
import type { LimitKey } from "@/lib/subscription/config";
import {
  checkAndConsume,
  LimitExceededError,
} from "@/lib/subscription/service";

export function withLimit<T extends Tool>(
  tool: T,
  kind: LimitKey,
  userId: string
): T {
  const originalExecute = tool.execute;
  if (!originalExecute) {
    return tool;
  }

  // biome-ignore lint/suspicious/noExplicitAny: wrapper has to be generic across all tool signatures
  const wrappedExecute = (async (input: any, options: any) => {
    try {
      await checkAndConsume(userId, kind);
    } catch (error) {
      if (error instanceof LimitExceededError) {
        return {
          error: "limit_exceeded",
          kind: error.kind,
          tier: error.tier,
          used: error.used,
          limit: error.limit,
          message: `You've hit your ${kind.replace("_", " ")} limit for today (${error.used}/${error.limit} on ${error.tier} tier). Upgrade to Pro for higher limits.`,
        };
      }
      throw error;
    }
    return originalExecute(input, options);
    // biome-ignore lint/suspicious/noExplicitAny: tool execute signature varies
  }) as any;

  return { ...tool, execute: wrappedExecute } as T;
}