vibn-agent-runner/vibn-frontend/lib/ai/gemini-client.ts

import type { LlmClient, StructuredCallArgs } from '@/lib/ai/llm-client';
import { GoogleGenAI } from '@google/genai';
import { zodToJsonSchema } from 'zod-to-json-schema';

// Use the new Google GenAI SDK (replacing the deprecated VertexAI SDK)
// Since Vertex AI gemini-3.1-pro-preview threw a 404 in your region, we use the standard AI Studio endpoint natively.
const DEFAULT_MODEL = process.env.GEMINI_MODEL || 'gemini-3.1-pro-preview';
const GOOGLE_API_KEY = process.env.GOOGLE_API_KEY || '';

if (!GOOGLE_API_KEY) {
  console.warn(`[GeminiLlmClient] WARNING: GOOGLE_API_KEY is not set. API calls will fail with a 403 Forbidden.`);
}

const ai = new GoogleGenAI({ apiKey: GOOGLE_API_KEY });

class JsonValidationError extends Error {
  constructor(message: string, public readonly rawResponse: string) {
    super(message);
  }
}

function extractJsonPayload(raw: string): string {
  const trimmed = raw.trim();
  if (trimmed.startsWith('```')) {
    return trimmed.replace(/^```(?:json)?/i, '').replace(/```$/, '').trim();
  }
  return trimmed;
}

export class GeminiLlmClient implements LlmClient {
  private readonly model: string;

  constructor() {
    this.model = DEFAULT_MODEL;
    console.log(`[GoogleGenAI] Initialized — model: ${this.model}`);
  }

  async structuredCall<TOutput>(args: StructuredCallArgs<TOutput>): Promise<TOutput> {
    if (args.model !== 'gemini') {
      throw new Error(`GeminiLlmClient only supports model "gemini" (got ${args.model})`);
    }

    const rawJsonSchema = zodToJsonSchema(args.schema, 'responseSchema') as any;
    let actualSchema: any = rawJsonSchema;
    if (rawJsonSchema.$ref && rawJsonSchema.definitions) {
      const refName = rawJsonSchema.$ref.replace('#/definitions/', '');
      actualSchema = rawJsonSchema.definitions[refName];
    }

    const convertToGoogleSchema = (schema: any): any => {
      if (!schema || typeof schema !== 'object') return schema;
      const out: any = {};
      if (schema.type) out.type = schema.type.toUpperCase();
      if (schema.properties) {
        out.properties = {};
        for (const [k, v] of Object.entries(schema.properties)) {
          out.properties[k] = convertToGoogleSchema(v);
        }
      }
      if (schema.items) out.items = convertToGoogleSchema(schema.items);
      if (schema.required) out.required = schema.required;
      if (schema.description) out.description = schema.description;
      if (schema.enum) out.enum = schema.enum;
      return out;
    };

    const googleSchema = convertToGoogleSchema(actualSchema);

    const contents = args.messages.map((m) => ({
      role: m.role === 'assistant' ? 'model' : 'user',
      parts: [{ text: m.content }],
    }));

    const config: any = {
      temperature: args.temperature ?? 1.0,
      responseMimeType: 'application/json',
      responseSchema: googleSchema,
      maxOutputTokens: 8192,
    };

    if (args.systemPrompt) {
      const exampleJson: any = {};
      for (const key of Object.keys(googleSchema.properties || {})) {
        exampleJson[key] = key === 'reply' ? 'Your response here' : null;
      }
      config.systemInstruction = `${args.systemPrompt}\n\nIMPERATIVE: Respond ONLY with this exact JSON format:\n${JSON.stringify(exampleJson)}\n\nDo NOT add any other fields.`;
    }

    if (args.thinking_config) {
      config.thinkingConfig = {
        thinkingBudget: 1024,
      };
    }

    const run = async (currentContents: any[]) => {
      console.log(`[GoogleGenAI] generateContent with ${this.model}`);
      
      const response = await ai.models.generateContent({
        model: this.model,
        contents: currentContents,
        config,
      });

      const text = response.text || '';

      if (text.trim().startsWith('<!DOCTYPE') || text.trim().startsWith('<html')) {
        throw new Error('GoogleGenAI returned HTML. Check API key permissions.');
      }
      if (!text) {
        throw new Error('Empty response from GoogleGenAI');
      }

      console.log('[GoogleGenAI] Response preview:', text.slice(0, 200));

      const cleaned = extractJsonPayload(text);

      let parsed: unknown;
      try {
        parsed = JSON.parse(cleaned);
      } catch (error) {
        throw new JsonValidationError(
          `Failed to parse JSON: ${(error as Error).message}`,
          text,
        );
      }

      const validation = args.schema.safeParse(parsed);
      if (!validation.success) {
        console.error('[GoogleGenAI] Schema validation failed:', validation.error.errors);
        throw new JsonValidationError(validation.error.message, text);
      }

      return validation.data;
    };

    try {
      return await run(contents);
    } catch (error) {
      if (!(error instanceof JsonValidationError)) throw error;

      console.warn(`[GoogleGenAI] JSON Validation failed. Retrying...`);
      const retryContents = [
        ...contents,
        {
          role: 'user',
          parts: [{
            text: `Your previous response was not valid JSON. Error: ${error.message}\nRespond again with ONLY valid JSON matching the schema. No code fences or comments.`,
          }],
        },
      ];
      return run(retryContents);
    }
  }
}