package gemini

type ChatRequest struct {
	Contents          []ChatContent        `json:"contents"`
	SafetySettings    []ChatSafetySettings `json:"safety_settings,omitempty"`
	GenerationConfig  ChatGenerationConfig `json:"generation_config,omitempty"`
	Tools             []ChatTools          `json:"tools,omitempty"`
	SystemInstruction *ChatContent         `json:"system_instruction,omitempty"`
}

type EmbeddingRequest struct {
	Model                string      `json:"model"`
	Content              ChatContent `json:"content"`
	TaskType             string      `json:"taskType,omitempty"`
	Title                string      `json:"title,omitempty"`
	OutputDimensionality int         `json:"outputDimensionality,omitempty"`
}

type BatchEmbeddingRequest struct {
	Requests []EmbeddingRequest `json:"requests"`
}

type EmbeddingData struct {
	Values []float64 `json:"values"`
}

type EmbeddingResponse struct {
	Embeddings []EmbeddingData `json:"embeddings"`
	Error      *Error          `json:"error,omitempty"`
}

type Error struct {
	Code    int    `json:"code,omitempty"`
	Message string `json:"message,omitempty"`
	Status  string `json:"status,omitempty"`
}

type InlineData struct {
	MimeType string `json:"mimeType"`
	Data     string `json:"data"`
}

type FunctionCall struct {
	FunctionName string `json:"name"`
	Arguments    any    `json:"args"`
}

type FunctionResponse struct {
	Name     string `json:"name"`
	Response any    `json:"response"`
}

type Part struct {
	Text             string            `json:"text,omitempty"`
	InlineData       *InlineData       `json:"inlineData,omitempty"`
	FunctionCall     *FunctionCall     `json:"functionCall,omitempty"`
	FunctionResponse *FunctionResponse `json:"functionResponse,omitempty"`
	// Thought marks this part as internal reasoning/thinking content (Gemini thinking models)
	Thought bool `json:"thought,omitempty"`
}

type ChatContent struct {
	Role  string `json:"role,omitempty"`
	Parts []Part `json:"parts"`
}

type ChatSafetySettings struct {
	Category  string `json:"category"`
	Threshold string `json:"threshold"`
}

type ChatTools struct {
	FunctionDeclarations any `json:"function_declarations,omitempty"`
}

type GeminiThinkingConfig struct {
	ThinkingBudget int `json:"thinkingBudget"` // -1 = dynamic, 0 = disabled
}

type ChatGenerationConfig struct {
	ResponseMimeType    string               `json:"responseMimeType,omitempty"`
	ResponseSchema      any                  `json:"responseSchema,omitempty"`
	ResponseModalities  []string             `json:"responseModalities,omitempty"`
	Temperature         *float64             `json:"temperature,omitempty"`
	TopP                *float64             `json:"topP,omitempty"`
	TopK                float64              `json:"topK,omitempty"`
	MaxOutputTokens     int                  `json:"maxOutputTokens,omitempty"`
	CandidateCount      int                  `json:"candidateCount,omitempty"`
	StopSequences       []string             `json:"stopSequences,omitempty"`
	ThinkingConfig      *GeminiThinkingConfig `json:"thinkingConfig,omitempty"`
}

// --- Google Imagen ---

// ImagenRequest is the request body for legacy Imagen (predict endpoint).
// POST /v1beta/models/{model}:predict
type ImagenRequest struct {
	Instances  []ImagenInstance `json:"instances"`
	Parameters ImagenParameters `json:"parameters"`
}

type ImagenInstance struct {
	Prompt string `json:"prompt"`
}

type ImagenParameters struct {
	SampleCount int    `json:"sampleCount"`
	AspectRatio string `json:"aspectRatio,omitempty"`
}

// GenerateImagesRequest is the request body for Imagen 3+ (generateImages endpoint).
// POST /v1/models/{model}:generateImages
type GenerateImagesRequest struct {
	Prompt         string `json:"prompt"`
	NumberOfImages int    `json:"number_of_images"`
	AspectRatio    string `json:"aspectRatio,omitempty"`
}

// EditImageRequest is the request body for Vertex AI image editing via :predict.
// Same endpoint as generation, uses instances[] with referenceImages. No parameters field.
type EditImageRequest struct {
	Instances []EditImageInstance `json:"instances"`
}

type EditImageInstance struct {
	Prompt          string               `json:"prompt"`
	ReferenceImages []ReferenceImageItem `json:"referenceImages"`
}

type ReferenceImageItem struct {
	ReferenceType  string             `json:"referenceType"`
	ReferenceId    int                `json:"referenceId"`
	ReferenceImage ReferenceImageData `json:"referenceImage"`
}

type ReferenceImageData struct {
	BytesBase64Encoded string `json:"bytesBase64Encoded"`
}

// ImagenResponse is the response from legacy Imagen :predict endpoint.
type ImagenResponse struct {
	Predictions []ImagenPrediction `json:"predictions"`
	Error       *Error             `json:"error,omitempty"`
}

// GenerateImagesResponse is the response from :generateImages and :editImage endpoints.
type GenerateImagesResponse struct {
	GeneratedImages []GeneratedImageItem `json:"generatedImages"`
	Error           *Error               `json:"error,omitempty"`
}

type GeneratedImageItem struct {
	Image    GeneratedImageData `json:"image"`
	MimeType string             `json:"mimeType"`
}

type GeneratedImageData struct {
	// ImageBytes is base64-encoded image data returned by :generateImages / :editImage
	ImageBytes string `json:"imageBytes"`
}

type ImagenPrediction struct {
	BytesBase64Encoded string `json:"bytesBase64Encoded"`
	MimeType           string `json:"mimeType"`
	GcsUri             string `json:"gcsUri"`
	Prompt             string `json:"prompt"`
}