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

const VERTEX_PROJECT_ID = process.env.VERTEX_AI_PROJECT_ID || 'gen-lang-client-0980079410';
const VERTEX_LOCATION = process.env.VERTEX_AI_LOCATION || 'global'; 
const DEFAULT_MODEL = process.env.VERTEX_AI_MODEL || 'gemini-2.0-flash-exp'; // Fast model for collector mode

// Helper to set up Google Application Credentials
function setupGoogleCredentials() {
  console.log('[Gemini Client] setupGoogleCredentials called');
  console.log('[Gemini Client] FIREBASE_CLIENT_EMAIL:', process.env.FIREBASE_CLIENT_EMAIL ? 'SET' : 'NOT SET');
  console.log('[Gemini Client] FIREBASE_PRIVATE_KEY:', process.env.FIREBASE_PRIVATE_KEY ? 'SET' : 'NOT SET');
  console.log('[Gemini Client] GOOGLE_APPLICATION_CREDENTIALS before:', process.env.GOOGLE_APPLICATION_CREDENTIALS || 'NOT SET');
  
  // Only set up if we have Firebase credentials and Google creds aren't already set
  if (process.env.FIREBASE_CLIENT_EMAIL && process.env.FIREBASE_PRIVATE_KEY && !process.env.GOOGLE_APPLICATION_CREDENTIALS) {
    const credentials = {
      type: 'service_account',
      project_id: VERTEX_PROJECT_ID,
      private_key_id: 'firebase-key',
      private_key: process.env.FIREBASE_PRIVATE_KEY.replace(/\\n/g, '\n'),
      client_email: process.env.FIREBASE_CLIENT_EMAIL,
      client_id: '',
      auth_uri: 'https://accounts.google.com/o/oauth2/auth',
      token_uri: 'https://oauth2.googleapis.com/token',
      auth_provider_x509_cert_url: 'https://www.googleapis.com/oauth2/v1/certs',
      client_x509_cert_url: `https://www.googleapis.com/robot/v1/metadata/x509/${encodeURIComponent(process.env.FIREBASE_CLIENT_EMAIL)}`,
      universe_domain: 'googleapis.com',
    };
    
    // Write credentials to a temp file that Google Auth Library can read
    const fs = require('fs');
    const os = require('os');
    const path = require('path');
    
    const tmpDir = os.tmpdir();
    const credPath = path.join(tmpDir, 'google-credentials.json');
    
    try {
      fs.writeFileSync(credPath, JSON.stringify(credentials));
      process.env.GOOGLE_APPLICATION_CREDENTIALS = credPath;
      console.log('[Gemini Client] ✅ Created credentials file at:', credPath);
      return true;
    } catch (error) {
      console.error('[Gemini Client] ❌ Failed to write credentials file:', error);
      return false;
    }
  } else {
    console.log('[Gemini Client] Skipping credentials setup - already set or missing Firebase creds');
  }
  return false;
}

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('```')) {
    const withoutFence = trimmed.replace(/^```(?:json)?/i, '').replace(/```$/, '');
    return withoutFence.trim();
  }
  return trimmed;
}

async function parseResponse<TOutput>(
  rawResponse: any,
  schema: StructuredCallArgs<TOutput>['schema'],
): Promise<TOutput> {
  // Extract text from Google GenAI response
  // The response structure is: { candidates: [{ content: { parts: [{ text: "..." }] } }] }
  let text = '';
  
  // Check for truncation
  const finishReason = rawResponse?.candidates?.[0]?.finishReason;
  if (finishReason && finishReason !== 'STOP') {
    console.warn(`[Gemini Client] 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;
  }
  
  // Check if we got HTML instead of JSON (API error)
  if (text.trim().startsWith('<!DOCTYPE') || text.trim().startsWith('<html')) {
    console.error('[Gemini Client] Received HTML instead of JSON. This usually means an API authentication or permission error.');
    console.error('[Gemini Client] Response preview:', text.substring(0, 500));
    throw new Error('Gemini API returned HTML instead of JSON. Check API permissions and authentication. See server logs for details.');
  }

  if (!text) {
    console.error('[Gemini Client] Empty response from API');
    console.error('[Gemini Client] Raw response:', JSON.stringify(rawResponse, null, 2)?.substring(0, 500));
    throw new Error('Empty response from Gemini API');
  }

  // Debug: Log what we received
  console.log('[Gemini Client] Received text:', text.substring(0, 300));

  const cleaned = extractJsonPayload(text);

  let parsed: unknown;
  try {
    parsed = JSON.parse(cleaned);
  } catch (error) {
    console.error('[Gemini Client] Failed to parse response as JSON');
    console.error('[Gemini Client] Raw text:', text.substring(0, 500));
    throw new JsonValidationError(
      `Failed to parse JSON response: ${(error as Error).message}`,
      text,
    );
  }

  // Debug: Log what we parsed
  console.log('[Gemini Client] Parsed JSON:', JSON.stringify(parsed, null, 2).substring(0, 300));

  const validation = schema.safeParse(parsed);
  if (!validation.success) {
    console.error('[Gemini Client] Schema validation failed!');
    console.error('[Gemini Client] Received JSON had these keys:', Object.keys(parsed as any));
    console.error('[Gemini Client] Validation errors:', validation.error.errors);
    throw new JsonValidationError(validation.error.message, text);
  }

  return validation.data;
}

export class GeminiLlmClient implements LlmClient {
  private readonly genAI: GoogleGenAI;
  private readonly model: string;
  private readonly location: string;
  private readonly projectId: string;

  constructor() {
    // Google GenAI SDK with Vertex AI support
    this.projectId = VERTEX_PROJECT_ID;
    this.location = VERTEX_LOCATION;
    this.model = DEFAULT_MODEL;
    
    // Set up Google Application Credentials BEFORE initializing the SDK
    setupGoogleCredentials();
    
    // Debug: Check environment variables
    console.log('[Gemini Client] Environment check:');
    console.log('  VERTEX_AI_PROJECT_ID:', process.env.VERTEX_AI_PROJECT_ID);
    console.log('  VERTEX_AI_LOCATION:', process.env.VERTEX_AI_LOCATION);
    console.log('  VERTEX_AI_MODEL:', process.env.VERTEX_AI_MODEL);
    console.log('  GOOGLE_APPLICATION_CREDENTIALS:', process.env.GOOGLE_APPLICATION_CREDENTIALS ? 'SET' : 'NOT SET');
    console.log('  FIREBASE_CLIENT_EMAIL:', process.env.FIREBASE_CLIENT_EMAIL ? 'SET' : 'NOT SET');
    
    // Initialize with Vertex AI configuration
    // The SDK will automatically use GOOGLE_APPLICATION_CREDENTIALS if set
    this.genAI = new GoogleGenAI({
      project: this.projectId,
      location: this.location,
      vertexai: true, // Enable Vertex AI mode
    });

    console.log(`[Gemini Client] Initialized with model: ${this.model}, location: ${this.location}`);
  }

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

    // Convert Zod schema to JSON Schema for Gemini
    const rawJsonSchema = zodToJsonSchema(args.schema, 'responseSchema');
    
    // Extract the actual schema from the definitions (zodToJsonSchema wraps it in $ref)
    let actualSchema = rawJsonSchema;
    const rawSchema = rawJsonSchema as any; // Type assertion for $ref access
    if (rawSchema.$ref && rawSchema.definitions) {
      const refName = rawSchema.$ref.replace('#/definitions/', '');
      actualSchema = rawSchema.definitions[refName];
    }
    
    // Convert to Google's expected format (UPPERCASE types)
    const convertToGoogleSchema = (schema: any): any => {
      if (!schema || typeof schema !== 'object') return schema;
      
      const converted: any = {};
      
      if (schema.type) {
        converted.type = schema.type.toUpperCase();
      }
      
      if (schema.properties) {
        converted.properties = {};
        for (const [key, value] of Object.entries(schema.properties)) {
          converted.properties[key] = convertToGoogleSchema(value);
        }
      }
      
      if (schema.items) {
        converted.items = convertToGoogleSchema(schema.items);
      }
      
      if (schema.required) {
        converted.required = schema.required;
      }
      
      if (schema.description) {
        converted.description = schema.description;
      }
      
      if (schema.enum) {
        converted.enum = schema.enum;
      }
      
      // Remove additionalProperties since Gemini doesn't use it
      // (it's a JSON Schema Draft 7 thing)
      
      return converted;
    };
    
    const googleSchema = convertToGoogleSchema(actualSchema);
    
    // Debug: Log the schema being sent
    console.log('[Gemini Client] Sending schema:', JSON.stringify(googleSchema, null, 2));
    
    // Build generation config matching Google's example structure
    const generationConfig: any = {
      temperature: args.temperature ?? 1.0,
      responseMimeType: 'application/json',
      responseSchema: googleSchema,
      maxOutputTokens: 32768, // Gemini 3 Pro supports up to 32k output tokens
    };
    
    // Main request object for REST API (flat structure)
    const config: any = {
      contents: [], // Will be populated below
      generationConfig: generationConfig,
    };

    // Add system instruction if provided
    if (args.systemPrompt) {
      // Create a minimal example showing the exact format
      const exampleJson: any = {};
      for (const [key, prop] of Object.entries(googleSchema.properties || {})) {
        if (key === 'reply') {
          exampleJson[key] = 'Your response here';
        } else {
          exampleJson[key] = null; // optional field
        }
      }
      
      config.systemInstruction = {
        parts: [{ 
          text: `${args.systemPrompt}\n\nIMPERATIVE: Respond ONLY with this exact JSON format:\n${JSON.stringify(exampleJson)}\n\nDo NOT add thought_process, response, or any other fields. Use only the keys shown above.`
        }],
      };
    }

    // Add thinking config if provided (for Gemini 3 Pro Preview)
    if (args.thinking_config) {
      config.generationConfig.thinkingConfig = {
        thinkingLevel: args.thinking_config.thinking_level?.toUpperCase() || 'HIGH',
        includeThoughts: args.thinking_config.include_thoughts || false,
      };
    }

    // Convert messages to Google GenAI format
    config.contents = args.messages.map((message) => ({
      role: message.role === 'assistant' ? 'model' : 'user',
      parts: [{ text: message.content }],
    }));

    const run = async () => {
      try {
        console.log('[Gemini Client] Calling generateContent via REST API...');
        
        // Use direct REST API call instead of SDK (SDK has auth issues)
        const { GoogleAuth } = require('google-auth-library');
        const auth = new GoogleAuth({
          scopes: ['https://www.googleapis.com/auth/cloud-platform'],
        });
        
        const client = await auth.getClient();
        const accessToken = await client.getAccessToken();
        
        const url = `https://aiplatform.googleapis.com/v1/projects/${this.projectId}/locations/${this.location}/publishers/google/models/${this.model}:generateContent`;
        
        console.log('[Gemini Client] Making request to:', url);
        
        const response = await fetch(url, {
          method: 'POST',
          headers: {
            'Authorization': `Bearer ${accessToken.token}`,
            'Content-Type': 'application/json',
          },
          body: JSON.stringify(config),
          signal: AbortSignal.timeout(180000), // 3 minute timeout
        });
        
        if (!response.ok) {
          const errorText = await response.text();
          throw new Error(`Vertex AI API error: ${response.status} ${errorText}`);
        }
        
        const result = await response.json();
        console.log('[Gemini Client] Got response from Gemini');
        
        return parseResponse(result, args.schema);
      } catch (error: any) {
        console.error('[Gemini Client] API call failed:', error.message || error);
        throw error;
      }
    };

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

      // Retry with error message
      config.contents = [
        ...config.contents,
        {
          role: 'user' as const,
          parts: [
            {
              text: `Your previous response was not valid JSON. Error: ${error.message}\n` +
                'Respond again with ONLY valid JSON that strictly matches the requested schema. Do not include comments or code fences.',
            },
          ],
        },
      ];

      return run();
    }
  }
}