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

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

const DEFAULT_MODEL = process.env.GEMINI_MODEL || 'gemini-2.0-flash-exp';
const GOOGLE_API_KEY = process.env.GOOGLE_API_KEY || '';

const GEMINI_BASE_URL = 'https://generativelanguage.googleapis.com/v1beta/models';

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;
}

async function parseResponse<TOutput>(
  rawResponse: any,
  schema: StructuredCallArgs<TOutput>['schema'],
): Promise<TOutput> {
  let text = '';

  const finishReason = rawResponse?.candidates?.[0]?.finishReason;
  if (finishReason && finishReason !== 'STOP') {
    console.warn(`[Gemini] WARNING: Response may be incomplete. finishReason: ${finishReason}`);
  }

  if (rawResponse?.candidates?.[0]?.content?.parts?.[0]?.text) {
    text = rawResponse.candidates[0].content.parts[0].text;
  } else if (rawResponse?.text) {
    text = rawResponse.text;
  } else if (typeof rawResponse === 'string') {
    text = rawResponse;
  }

  if (text.trim().startsWith('<!DOCTYPE') || text.trim().startsWith('<html')) {
    console.error('[Gemini] Received HTML — likely an API auth error');
    throw new Error('Gemini API returned HTML. Check GOOGLE_API_KEY.');
  }

  if (!text) {
    console.error('[Gemini] Empty response:', JSON.stringify(rawResponse)?.slice(0, 300));
    throw new Error('Empty response from Gemini API');
  }

  console.log('[Gemini] 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 = schema.safeParse(parsed);
  if (!validation.success) {
    console.error('[Gemini] Schema validation failed:', validation.error.errors);
    throw new JsonValidationError(validation.error.message, text);
  }

  return validation.data;
}

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

  constructor() {
    this.model = DEFAULT_MODEL;
    if (!GOOGLE_API_KEY) {
      console.warn('[Gemini] WARNING: GOOGLE_API_KEY is not set');
    }
    console.log(`[Gemini] 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})`);
    }

    // Convert Zod schema → Google schema format
    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);

    // Build request body
    const body: any = {
      contents: args.messages.map((m) => ({
        role: m.role === 'assistant' ? 'model' : 'user',
        parts: [{ text: m.content }],
      })),
      generationConfig: {
        temperature: args.temperature ?? 1.0,
        responseMimeType: 'application/json',
        responseSchema: googleSchema,
        maxOutputTokens: 32768,
      },
    };

    if (args.systemPrompt) {
      const exampleJson: any = {};
      for (const key of Object.keys(googleSchema.properties || {})) {
        exampleJson[key] = key === 'reply' ? 'Your response here' : null;
      }
      body.systemInstruction = {
        parts: [{
          text: `${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) {
      body.generationConfig.thinkingConfig = {
        thinkingLevel: args.thinking_config.thinking_level?.toUpperCase() || 'HIGH',
        includeThoughts: args.thinking_config.include_thoughts || false,
      };
    }

    const url = `${GEMINI_BASE_URL}/${this.model}:generateContent?key=${GOOGLE_API_KEY}`;

    const run = async () => {
      console.log(`[Gemini] POST ${GEMINI_BASE_URL}/${this.model}:generateContent`);
      const response = await fetch(url, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(body),
        signal: AbortSignal.timeout(180_000),
      });

      if (!response.ok) {
        const errorText = await response.text();
        throw new Error(`Gemini API error (${response.status}): ${errorText}`);
      }

      const result = await response.json();
      return parseResponse(result, args.schema);
    };

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

      // Retry once on JSON parse failure
      body.contents = [
        ...body.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();
    }
  }
}