import { GoogleGenerativeAI, GenerativeModel, GenerationConfig } from '@google/generative-ai';
// Alternative: import { genai } from 'google-genai'; // For the new Google GenAI SDK

export interface GeminiModelConfig {
  model: string;
  maxTokens?: number;
  temperature?: number;
  topP?: number;
  topK?: number;
  stopSequences?: string[];
}

export interface GeminiResponse {
  text: string;
  finishReason?: string;
  usage?: {
    promptTokens: number;
    completionTokens: number;
    totalTokens: number;
  };
}

export interface GeminiStreamChunk {
  text: string;
  done: boolean;
  finishReason?: string;
}

class GeminiClient {
  private client: GoogleGenerativeAI | null = null;
  private models: Map<string, GenerativeModel> = new Map();

  constructor() {
    // Lazy initialization - client will be created when first needed
  }

  private getClient(): GoogleGenerativeAI {
    if (!this.client) {
      const apiKey = process.env.GEMINI_API_KEY;
      if (!apiKey) {
        throw new Error('GEMINI_API_KEY environment variable is required');
      }
      this.client = new GoogleGenerativeAI(apiKey);
    }
    return this.client;
  }

  /**
   * Get or create a generative model instance
   */
  private getModel(modelName: string, config?: Partial<GenerationConfig>): GenerativeModel {
    const cacheKey = `${modelName}-${JSON.stringify(config || {})}`;
    
    if (!this.models.has(cacheKey)) {
      const generationConfig: GenerationConfig = {
        temperature: config?.temperature ?? 0.7,
        topP: config?.topP ?? 0.8,
        topK: config?.topK ?? 40,
        maxOutputTokens: config?.maxOutputTokens ?? 8192,
        stopSequences: config?.stopSequences ?? [],
      };

      const model = this.getClient().getGenerativeModel({
        model: modelName,
        generationConfig,
      });

      this.models.set(cacheKey, model);
    }

    return this.models.get(cacheKey)!;
  }

  /**
   * Generate content using Gemini models
   */
  async generateContent(
    modelName: string,
    prompt: string,
    config?: Partial<GeminiModelConfig>
  ): Promise<GeminiResponse> {
    try {
      const model = this.getModel(modelName, {
        temperature: config?.temperature,
        topP: config?.topP,
        topK: config?.topK,
        maxOutputTokens: config?.maxTokens,
        stopSequences: config?.stopSequences,
      });

      const result = await model.generateContent(prompt);
      const response = result.response;
      
      return {
        text: response.text(),
        finishReason: response.candidates?.[0]?.finishReason,
        usage: {
          promptTokens: response.usageMetadata?.promptTokenCount || 0,
          completionTokens: response.usageMetadata?.candidatesTokenCount || 0,
          totalTokens: response.usageMetadata?.totalTokenCount || 0,
        }
      };
    } catch (error) {
      console.error('Error generating content with Gemini:', error);
      throw new Error(`Gemini generation failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  /**
   * Generate content stream using Gemini models
   */
  async* generateContentStream(
    modelName: string,
    prompt: string,
    config?: Partial<GeminiModelConfig>
  ): AsyncGenerator<GeminiStreamChunk, void, unknown> {
    try {
      const model = this.getModel(modelName, {
        temperature: config?.temperature,
        topP: config?.topP,
        topK: config?.topK,
        maxOutputTokens: config?.maxTokens,
        stopSequences: config?.stopSequences,
      });

      const result = await model.generateContentStream(prompt);
      
      for await (const chunk of result.stream) {
        const chunkText = chunk.text();
        yield {
          text: chunkText,
          done: false,
          finishReason: chunk.candidates?.[0]?.finishReason,
        };
      }

      // Final chunk to indicate completion
      yield {
        text: '',
        done: true,
      };
    } catch (error) {
      console.error('Error generating content stream with Gemini:', error);
      throw new Error(`Gemini stream generation failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  /**
   * Generate content with system instruction and user prompt
   */
  async generateWithSystemInstruction(
    modelName: string,
    systemInstruction: string,
    userPrompt: string,
    config?: Partial<GeminiModelConfig>
  ): Promise<GeminiResponse> {
    const fullPrompt = `${systemInstruction}\n\nUser: ${userPrompt}\n\nAssistant:`;
    return this.generateContent(modelName, fullPrompt, config);
  }

  /**
   * Generate content with conversation history
   */
  async generateWithHistory(
    modelName: string,
    messages: Array<{ role: 'user' | 'assistant'; content: string }>,
    config?: Partial<GeminiModelConfig>
  ): Promise<GeminiResponse> {
    try {
      const model = this.getModel(modelName, {
        temperature: config?.temperature,
        topP: config?.topP,
        topK: config?.topK,
        maxOutputTokens: config?.maxTokens,
        stopSequences: config?.stopSequences,
      });

      // Convert messages to Gemini format
      const history = messages.slice(0, -1).map(msg => ({
        role: msg.role === 'user' ? 'user' : 'model',
        parts: [{ text: msg.content }],
      }));

      const lastMessage = messages[messages.length - 1];
      
      const chat = model.startChat({
        history: history as Array<{ role: 'user' | 'model'; parts: Array<{ text: string }> }>,
      });

      const result = await chat.sendMessage(lastMessage.content);
      const response = result.response;

      return {
        text: response.text(),
        finishReason: response.candidates?.[0]?.finishReason,
        usage: {
          promptTokens: response.usageMetadata?.promptTokenCount || 0,
          completionTokens: response.usageMetadata?.candidatesTokenCount || 0,
          totalTokens: response.usageMetadata?.totalTokenCount || 0,
        }
      };
    } catch (error) {
      console.error('Error generating content with history:', error);
      throw new Error(`Gemini chat generation failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  /**
   * Count tokens in text
   */
  async countTokens(modelName: string, text: string): Promise<number> {
    try {
      const model = this.getModel(modelName);
      const result = await model.countTokens(text);
      return result.totalTokens;
    } catch (error) {
      console.error('Error counting tokens:', error);
      // Fallback estimation: roughly 4 characters per token
      return Math.ceil(text.length / 4);
    }
  }

  /**
   * Validate model availability
   */
  async validateModel(modelName: string): Promise<boolean> {
    try {
      const model = this.getModel(modelName);
      // Try a simple generation to test the model
      await model.generateContent('Test');
      return true;
    } catch (error) {
      console.error(`Model ${modelName} validation failed:`, error);
      return false;
    }
  }

  /**
   * Get available models (static list for now)
   */
  getAvailableModels(): string[] {
    return [
      'gemini-2.5-flash',
      'gemini-2.5-pro',
      'gemini-2.5-flash-lite',
      'gemma-3-27b-it',
      'gemini-2.0-flash'
    ];
  }

  /**
   * Clear model cache
   */
  clearCache(): void {
    this.models.clear();
  }

  /**
   * Get model info
   */
  getModelInfo(modelName: string): {
    name: string;
    maxTokens: number;
    features: string[];
    description: string;
  } {
    const modelInfo: Record<string, { name: string; maxTokens: number; features: string[]; description: string }> = {
      'gemini-2.5-flash': {
        name: 'Gemini 2.5 Flash',
        maxTokens: 1_048_576,
        features: ['text', 'code', 'multimodal', 'fast'],
        description: 'Fast and efficient model for general tasks'
      },
      'gemini-2.5-pro': {
        name: 'Gemini 2.5 Pro',
        maxTokens: 2_097_152,
        features: ['text', 'code', 'multimodal', 'reasoning', 'advanced'],
        description: 'Most capable model for complex reasoning tasks'
      },
      'gemini-2.5-flash-lite': {
        name: 'Gemini 2.5 Flash Lite',
        maxTokens: 1_048_576,
        features: ['text', 'code', 'fast', 'lightweight'],
        description: 'Lightweight version optimized for speed'
      },
      'gemma-3-27b-it': {
        name: 'Gemma 3 27B IT',
        maxTokens: 8_192,
        features: ['text', 'code', 'instruction', 'specialized'],
        description: 'Instruction-tuned model for specific tasks'
      },
      'gemini-2.0-flash': {
        name: 'Gemini 2.0 Flash',
        maxTokens: 1_048_576,
        features: ['text', 'code', 'multimodal', 'latest'],
        description: 'Latest generation model with improved capabilities'
      }
    };

    return modelInfo[modelName] || {
      name: modelName,
      maxTokens: 8192,
      features: ['text'],
      description: 'Unknown model'
    };
  }
}

// Export singleton instance
export const geminiClient = new GeminiClient();

// Export class for testing
export { GeminiClient };