#!/usr/bin/env python3
import google.generativeai as genai
import sys
import json
import base64

def log(msg):
    """Print to stderr for logging"""
    print(msg, file=sys.stderr, flush=True)

def query_gemini(api_key, prompt, image_path=None):
    log("\n=== Gemini API Request ===")
    log(f"API Key: {api_key[:4]}...{api_key[-4:]}")  # Show only first/last 4 chars
    log(f"Prompt:\n{prompt}")
    
    # Configure the Gemini API
    genai.configure(api_key=api_key)
    model = genai.GenerativeModel('gemini-2.0-flash-exp')
    
    try:
        if image_path:
            log(f"\nProcessing image: {image_path}")
            with open(image_path, 'rb') as f:
                image_data = f.read()
            log(f"Image size: {len(image_data):,} bytes")
            
            image_parts = [
                {
                    "mime_type": "image/jpeg",
                    "data": base64.b64encode(image_data).decode('utf-8')
                }
            ]
            response = model.generate_content([prompt, image_parts[0]])
        else:
            log("\nProcessing text-only request")
            response = model.generate_content(prompt)
        
        log("\n=== Gemini API Response ===")
        log(f"Response:\n{response.text}")
        
        # Create the JSON response
        result = {
            "success": True,
            "text": response.text
        }
        
        # Log the result to stderr
        log("\n=== JSON Response ===")
        log(json.dumps(result, indent=2))
        
        # Return the compact JSON to stdout without any extra newlines
        return json.dumps(result, separators=(',', ':'))
    except Exception as e:
        log("\n=== Gemini API Error ===")
        log(f"Error: {str(e)}")
        
        error_result = {
            "success": False,
            "error": str(e)
        }
        
        # Log the error to stderr
        log("\n=== JSON Error Response ===")
        log(json.dumps(error_result, indent=2))
        
        # Return the compact JSON to stdout without any extra newlines
        return json.dumps(error_result, separators=(',', ':'))

if __name__ == "__main__":
    if len(sys.argv) < 3:
        log("Usage: query_gemini_api.py <api_key> <prompt> [image_path]")
        sys.exit(1)
    
    api_key = sys.argv[1]
    prompt = sys.argv[2]
    image_path = sys.argv[3] if len(sys.argv) > 3 else None
    
    # Get the result and print it to stdout without any extra newlines
    result = query_gemini(api_key, prompt, image_path)
    print(result, end='', flush=True)