package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"

	"github.com/google/generative-ai-go/genai"
	"google.golang.org/api/option"
)

// Response represents the JSON response structure
type Response struct {
	Success bool   `json:"success"`
	Text    string `json:"text,omitempty"`
	Error   string `json:"error,omitempty"`
}

// logToStderr logs messages to stderr for debugging
func logToStderr(format string, a ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", a...)
}

func main() {
	if len(os.Args) < 3 {
		logToStderr("Usage: %s <api_key> <prompt> [image_path]", os.Args[0])
		os.Exit(1)
	}

	apiKey := os.Args[1]
	prompt := os.Args[2]
	var imagePath string
	if len(os.Args) > 3 {
		imagePath = os.Args[3]
	}

	logToStderr("\n=== Gemini API Request ===")
	logToStderr("API Key: %s...%s", apiKey[:4], apiKey[len(apiKey)-4:])
	logToStderr("Prompt:\n%s", prompt)

	ctx := context.Background()
	client, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
	if err != nil {
		handleError(fmt.Errorf("failed to create client: %v", err))
		return
	}
	defer client.Close()

	model := client.GenerativeModel("gemini-2.0-flash-exp")
	var response *genai.GenerateContentResponse

	if imagePath != "" {
		logToStderr("\nProcessing image: %s", imagePath)

		imageData, err := os.ReadFile(imagePath)
		if err != nil {
			handleError(fmt.Errorf("failed to read image: %v", err))
			return
		}
		logToStderr("Image size: %d bytes", len(imageData))

		model = client.GenerativeModel("gemini-2.0-flash-exp")
		response, err = model.GenerateContent(ctx, genai.Text(prompt), genai.ImageData("image/jpeg", imageData))
	} else {
		logToStderr("\nProcessing text-only request")
		response, err = model.GenerateContent(ctx, genai.Text(prompt))
	}

	if err != nil {
		handleError(fmt.Errorf("failed to generate content: %v", err))
		return
	}

	if len(response.Candidates) == 0 {
		handleError(fmt.Errorf("no response candidates received"))
		return
	}

	text := fmt.Sprintf("%v", response.Candidates[0].Content.Parts[0])
	logToStderr("\n=== Gemini API Response ===")
	logToStderr("Response:\n%s", text)

	result := Response{
		Success: true,
		Text:    text,
	}

	logToStderr("\n=== JSON Response ===")
	prettyJSON, _ := json.MarshalIndent(result, "", "  ")
	logToStderr(string(prettyJSON))

	// Output the final JSON to stdout without indentation
	finalJSON, _ := json.Marshal(result)
	fmt.Print(string(finalJSON))
}

func handleError(err error) {
	logToStderr("\n=== Gemini API Error ===")
	logToStderr("Error: %v", err)

	errorResult := Response{
		Success: false,
		Error:   err.Error(),
	}

	logToStderr("\n=== JSON Error Response ===")
	prettyJSON, _ := json.MarshalIndent(errorResult, "", "  ")
	logToStderr(string(prettyJSON))

	finalJSON, _ := json.Marshal(errorResult)
	fmt.Print(string(finalJSON))
}