import { createVertexProvider } from "@ai-sdk/google-vertex";
// Assuming other necessary imports

const provider = createVertexProvider({
  // ... other options
  fetch: async (input, init) => {
    const request = new Request(input, init);
    const headers = request.headers;
    const contentType = headers.get("content-type");

    // Check if it's a JSON request likely targeting the generateContent endpoint
    if (contentType?.includes("application/json")) {
      // Clone the request to read the body safely
      const clonedRequest = request.clone();
      try {
        const body = await clonedRequest.json<Record<string, unknown | undefined>>();

        // If cachedContent is present, remove fields disallowed by the API
        if (body?.cachedContent) {
          delete body.tools;
          delete body.toolConfig;
          delete body.systemInstruction;

          // Create a new init object with the modified body
          const newInit = {
            ...init,
            body: JSON.stringify(body), // Reserialize the modified body
          };
          // Fetch with the modified init object
          return fetch(input, newInit);
        }
      } catch (error) {
        console.error("Error processing request body:", error);
        // Proceed with the original request if JSON parsing fails or it's not the expected structure
      }
    }
    // Fallback to the original fetch call
    return fetch(input, init);
  },
});

import { GoogleGenerativeAI, FunctionCallingConfigMode, Schema } from "@google/generative-ai";
import { asSchema, convertJSONSchemaToOpenAPISchema } from "@ai-sdk/provider-utils";
import dayjs from "dayjs"; // Example date library

// Define tools using Vercel AI SDK's tool definition
const tools = {
  think: tool({ /* ... think tool definition */ }),
  search: tool({ /* ... search tool definition */ }),
};

// Define system instruction
const systemInstructionText = "You're a friendly bot with a very long system instruction and potentially large files referenced in the cache.";

// Prepare tool definitions for Google API format
const functionDeclarations = Object.entries(tools).map(([name, toolDef]) => ({
  name,
  description: toolDef.description ?? "",
  parameters: convertJSONSchemaToOpenAPISchema(
    "jsonSchema" in toolDef.parameters
      ? toolDef.parameters.jsonSchema
      : asSchema(toolDef.parameters as any).jsonSchema ?? {}
  ) as Schema,
}));

// ToolConfig for Google API
const toolConfig = {
  functionCallingConfig: {
    mode: FunctionCallingConfigMode.AUTO,
    // allowedFunctionNames: Object.keys(tools), // Only needed if Mode is ANY
  },
};

// Initialize Google SDK (ensure auth is configured)
const genai = new GoogleGenerativeAI({
  vertexai: true, // or false for Google AI
  project: "your-gcp-project-id",
  // googleAuthOptions: { ... }, // If needed
});

// Example conversation ID or unique identifier for the cache
const conversationId = "unique-conversation-123";

// Create the cache
const cache = await genai.caches.create({
  model: "gemini-1.5-flash-001", // Ensure this matches the model used later
  displayName: `Cache for conversation ${conversationId}`,
  ttl: "600s", // How long the cache is allowed to be valid
  systemInstruction: {
    parts: [{ text: systemInstructionText }],
  },
  tools: [{ functionDeclarations }],
  toolConfig: toolConfig,
  contents: [ // Include initial context/files here
    {
      role: "user", // Can be 'user' or 'model'
      parts: [
        { text: "Here is a large document for context:" },
        {
          fileData: {
            // Use GCS URI for large files
            fileUri: "gs://your-bucket-name/your-large-file.pdf",
            mimeType: "application/pdf",
            // Alternatively, for small data:
            // inlineData: { data: Buffer.from("...").toString("base64"), mimeType: "text/plain" }
          },
        },
      ],
    },
    // Add more history/context if needed
  ],
});

console.log(`Cache created: ${cache.name}`);

import { generateText } from "ai";
import { CoreMessage } from "ai";

// Use the modified provider and pass the cache name
const model = provider(cache.model); // Use the same model name as the cache

const messages: CoreMessage[] = [
  // Add only *new* messages not already in the cache's 'contents'
  { role: "user", content: "Summarize the document provided earlier." }
];

const result = await generateText({
  model: model,
  // Pass cache name via provider options defined in Step 1
  providerOptions: {
    cachedContent: cache.name,
  },
  system: systemInstructionText,
  tools: tools, // Still required here for executing functions
  messages: messages, // Only new messages
  maxSteps: 20, // Maximum amount of steps / tool calls
});

console.log(result.text);