Pulse/internal/ai/providers/gemini.go

package providers

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/rs/zerolog/log"
)

const (
	geminiAPIURL         = "https://generativelanguage.googleapis.com/v1beta"
	geminiMaxRetries     = 3
	geminiInitialBackoff = 2 * time.Second
)

// GeminiClient implements the Provider interface for Google's Gemini API
type GeminiClient struct {
	apiKey  string
	model   string
	baseURL string
	client  *http.Client
}

// NewGeminiClient creates a new Gemini API client
// timeout is optional - pass 0 to use the default 5 minute timeout
func NewGeminiClient(apiKey, model, baseURL string, timeout time.Duration) *GeminiClient {
	if baseURL == "" {
		baseURL = geminiAPIURL
	}
	// Strip provider prefix if present - the model should be just the model name
	model = strings.TrimPrefix(model, "gemini:")
	if timeout <= 0 {
		timeout = 300 * time.Second // Default 5 minutes
	}
	return &GeminiClient{
		apiKey:  apiKey,
		model:   model,
		baseURL: baseURL,
		client: &http.Client{
			Timeout: timeout,
		},
	}
}

// Name returns the provider name
func (c *GeminiClient) Name() string {
	return "gemini"
}

// geminiRequest is the request body for the Gemini API
type geminiRequest struct {
	Contents          []geminiContent         `json:"contents"`
	SystemInstruction *geminiContent          `json:"systemInstruction,omitempty"`
	GenerationConfig  *geminiGenerationConfig `json:"generationConfig,omitempty"`
	Tools             []geminiToolDef         `json:"tools,omitempty"`
	ToolConfig        *geminiToolConfig       `json:"toolConfig,omitempty"`
}

// geminiToolConfig controls how the model uses tools
// See: https://ai.google.dev/api/caching#ToolConfig
type geminiToolConfig struct {
	FunctionCallingConfig *geminiFunctionCallingConfig `json:"functionCallingConfig,omitempty"`
}

type geminiFunctionCallingConfig struct {
	Mode string `json:"mode"` // AUTO, ANY, or NONE
}

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

type geminiPart struct {
	Text                  string                  `json:"text,omitempty"`
	FunctionCall          *geminiFunctionCall     `json:"functionCall,omitempty"`
	FunctionResponse      *geminiFunctionResponse `json:"functionResponse,omitempty"`
	ThoughtSignature      json.RawMessage         `json:"thoughtSignature,omitempty"`
	ThoughtSignatureSnake json.RawMessage         `json:"thought_signature,omitempty"`
}

type geminiFunctionCall struct {
	ID   string                 `json:"id,omitempty"`
	Name string                 `json:"name"`
	Args map[string]interface{} `json:"args"`
}

type geminiFunctionResponse struct {
	ID       string `json:"id,omitempty"`
	Name     string `json:"name"`
	Response struct {
		Content string `json:"content"`
	} `json:"response"`
}

type geminiGenerationConfig struct {
	MaxOutputTokens int     `json:"maxOutputTokens,omitempty"`
	Temperature     float64 `json:"temperature,omitempty"`
}

type geminiToolDef struct {
	FunctionDeclarations []geminiFunctionDeclaration `json:"functionDeclarations,omitempty"`
}

type geminiFunctionDeclaration struct {
	Name        string                 `json:"name"`
	Description string                 `json:"description,omitempty"`
	Parameters  map[string]interface{} `json:"parameters,omitempty"`
}

// geminiResponse is the response from the Gemini API
type geminiResponse struct {
	Candidates     []geminiCandidate     `json:"candidates"`
	UsageMetadata  *geminiUsageMetadata  `json:"usageMetadata"`
	PromptFeedback *geminiPromptFeedback `json:"promptFeedback,omitempty"`
}

type geminiCandidate struct {
	Content       geminiContent  `json:"content"`
	FinishReason  string         `json:"finishReason"`
	SafetyRatings []geminySafety `json:"safetyRatings,omitempty"`
}

type geminySafety struct {
	Category    string `json:"category"`
	Probability string `json:"probability"`
	Blocked     bool   `json:"blocked"`
}

type geminiPromptFeedback struct {
	BlockReason   string         `json:"blockReason,omitempty"`
	SafetyRatings []geminySafety `json:"safetyRatings,omitempty"`
}

type geminiUsageMetadata struct {
	PromptTokenCount     int `json:"promptTokenCount"`
	CandidatesTokenCount int `json:"candidatesTokenCount"`
	TotalTokenCount      int `json:"totalTokenCount"`
}

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

// sanitizeGeminiContents validates and repairs message ordering for Gemini's constraints.
// Gemini requires that a model message containing function calls must be immediately
// followed by a user message containing function responses. If pruning or errors
// leave orphaned function calls (model+functionCalls not followed by function responses),
// this strips the function call parts, keeping only text content if present.
func sanitizeGeminiContents(contents []geminiContent) []geminiContent {
	result := make([]geminiContent, 0, len(contents))

	for i, c := range contents {
		// Check if this is a model message with function calls
		hasFunctionCall := false
		for _, p := range c.Parts {
			if p.FunctionCall != nil {
				hasFunctionCall = true
				break
			}
		}

		if c.Role == "model" && hasFunctionCall {
			// Check if the next message is a user message with function responses
			hasFollowingResponse := false
			if i+1 < len(contents) {
				next := contents[i+1]
				if next.Role == "user" {
					for _, p := range next.Parts {
						if p.FunctionResponse != nil {
							hasFollowingResponse = true
							break
						}
					}
				}
			}

			if !hasFollowingResponse {
				// Orphaned function calls — strip them, keep text only
				var textParts []geminiPart
				for _, p := range c.Parts {
					if p.Text != "" && p.FunctionCall == nil {
						textParts = append(textParts, geminiPart{Text: p.Text})
					}
				}
				if len(textParts) > 0 {
					result = append(result, geminiContent{
						Role:  c.Role,
						Parts: textParts,
					})
				}
				log.Debug().
					Int("message_index", i).
					Msg("[Gemini] Stripped orphaned function calls from model message")
				continue
			}
		}

		// Check if this is a user message with function responses
		// that isn't preceded by a model message with function calls
		hasFunctionResponse := false
		for _, p := range c.Parts {
			if p.FunctionResponse != nil {
				hasFunctionResponse = true
				break
			}
		}

		if c.Role == "user" && hasFunctionResponse {
			hasPrecedingCall := false
			if i > 0 {
				prev := contents[i-1]
				if prev.Role == "model" {
					for _, p := range prev.Parts {
						if p.FunctionCall != nil {
							hasPrecedingCall = true
							break
						}
					}
				}
			}
			// Also check if the preceding message in result has function calls
			// (it might have been the immediately previous content we just added)
			if !hasPrecedingCall && len(result) > 0 {
				prev := result[len(result)-1]
				if prev.Role == "model" {
					for _, p := range prev.Parts {
						if p.FunctionCall != nil {
							hasPrecedingCall = true
							break
						}
					}
				}
			}

			if !hasPrecedingCall {
				// Orphaned function responses — drop them
				log.Debug().
					Int("message_index", i).
					Msg("[Gemini] Dropped orphaned function responses from user message")
				continue
			}
		}

		result = append(result, c)
	}

	return result
}

// convertToolChoiceToGemini converts our ToolChoice to Gemini's mode string.
// Pulse omits automatic tool config so normal tool use stays model-owned.
// See: https://ai.google.dev/gemini-api/docs/generate-content#function_calling
func convertToolChoiceToGemini(tc *ToolChoice) string {
	if tc == nil {
		return ""
	}
	switch tc.Type {
	case ToolChoiceNone:
		return "NONE"
	case ToolChoiceRequired:
		return "ANY"
	default:
		return ""
	}
}

func buildGeminiFunctionDeclarations(tools []Tool) []geminiFunctionDeclaration {
	funcDecls := make([]geminiFunctionDeclaration, 0, len(tools))
	for _, t := range tools {
		if t.Type != "" && t.Type != "function" {
			continue
		}
		funcDecls = append(funcDecls, geminiFunctionDeclaration{
			Name:        t.Name,
			Description: t.Description,
			Parameters:  sanitizeGeminiToolSchema(t.InputSchema),
		})
	}
	return funcDecls
}

func applyGeminiTools(geminiReq *geminiRequest, tools []Tool, toolChoice *ToolChoice) []geminiFunctionDeclaration {
	shouldAddTools := len(tools) > 0
	if toolChoice != nil && toolChoice.Type == ToolChoiceNone {
		shouldAddTools = false
	}
	if !shouldAddTools {
		return nil
	}

	funcDecls := buildGeminiFunctionDeclarations(tools)
	if len(funcDecls) == 0 {
		return nil
	}
	geminiReq.Tools = []geminiToolDef{{FunctionDeclarations: funcDecls}}
	if mode := convertToolChoiceToGemini(toolChoice); mode != "" {
		geminiReq.ToolConfig = &geminiToolConfig{
			FunctionCallingConfig: &geminiFunctionCallingConfig{Mode: mode},
		}
	}
	return funcDecls
}

// sanitizeGeminiToolSchema projects Pulse's provider-neutral JSON schema onto
// Gemini's function declaration subset. Gemini rejects additionalProperties in
// function parameters, while other providers rely on it for stricter schemas.
func sanitizeGeminiToolSchema(schema map[string]interface{}) map[string]interface{} {
	if schema == nil {
		return nil
	}
	sanitized, _ := sanitizeGeminiSchemaValue(schema).(map[string]interface{})
	return sanitized
}

func sanitizeGeminiSchemaValue(value interface{}) interface{} {
	switch v := value.(type) {
	case map[string]interface{}:
		out := make(map[string]interface{}, len(v))
		for key, item := range v {
			if key == "additionalProperties" {
				continue
			}
			out[key] = sanitizeGeminiSchemaValue(item)
		}
		return out
	case []interface{}:
		out := make([]interface{}, len(v))
		for i, item := range v {
			out[i] = sanitizeGeminiSchemaValue(item)
		}
		return out
	case []string:
		return append([]string(nil), v...)
	default:
		return v
	}
}

// convertMessagesToGemini converts provider-neutral messages into Gemini's
// content shape used by both non-streaming and streaming requests. Gemini
// requires function responses to reference function names rather than provider
// call ids, and consecutive tool results are grouped into one user content
// block after the model function calls.
func convertMessagesToGemini(reqMessages []Message) []geminiContent {
	contents := make([]geminiContent, 0, len(reqMessages))

	for i := 0; i < len(reqMessages); i++ {
		m := reqMessages[i]
		if m.Role == "system" {
			continue
		}

		role := m.Role
		if role == "assistant" {
			role = "model"
		}

		if m.ToolResult != nil {
			var assistantMsg *Message
			if i > 0 && reqMessages[i-1].Role == "assistant" {
				assistantMsg = &reqMessages[i-1]
			}
			resolveName := func(id string) string {
				if assistantMsg != nil {
					for _, call := range assistantMsg.ToolCalls {
						if call.ID == id {
							return call.Name
						}
					}
				}
				return id
			}

			parts := []geminiPart{geminiFunctionResponsePart(m.ToolResult.ToolUseID, resolveName(m.ToolResult.ToolUseID), m.ToolResult.Content)}
			for i+1 < len(reqMessages) {
				next := reqMessages[i+1]
				if next.ToolResult == nil {
					break
				}
				parts = append(parts, geminiFunctionResponsePart(next.ToolResult.ToolUseID, resolveName(next.ToolResult.ToolUseID), next.ToolResult.Content))
				i++
			}

			contents = append(contents, geminiContent{
				Role:  "user",
				Parts: parts,
			})
			continue
		}

		if m.Role == "assistant" && len(m.ToolCalls) > 0 {
			parts := make([]geminiPart, 0, len(m.ToolCalls)+1)
			if m.Content != "" {
				parts = append(parts, geminiPart{Text: m.Content})
			}
			for _, tc := range m.ToolCalls {
				parts = append(parts, geminiPart{
					FunctionCall: &geminiFunctionCall{
						ID:   tc.ID,
						Name: tc.Name,
						Args: tc.Input,
					},
					ThoughtSignature: tc.ThoughtSignature,
				})
			}
			contents = append(contents, geminiContent{
				Role:  "model",
				Parts: parts,
			})
			continue
		}

		if m.Content == "" {
			continue
		}

		contents = append(contents, geminiContent{
			Role: role,
			Parts: []geminiPart{
				{Text: m.Content},
			},
		})
	}

	return sanitizeGeminiContents(contents)
}

func geminiFunctionResponsePart(id, name, content string) geminiPart {
	response := geminiFunctionResponse{ID: id, Name: name}
	response.Response.Content = content
	return geminiPart{FunctionResponse: &response}
}

// Chat sends a chat request to the Gemini API
func (c *GeminiClient) Chat(ctx context.Context, req ChatRequest) (*ChatResponse, error) {
	// Convert messages to Gemini format
	contents := convertMessagesToGemini(req.Messages)

	// Use provided model or fall back to client default
	model := req.Model
	// Strip provider prefix if present
	if strings.HasPrefix(model, "gemini:") {
		model = strings.TrimPrefix(model, "gemini:")
	}
	if model == "" {
		model = c.model
	}

	geminiReq := geminiRequest{
		Contents: contents,
	}

	// Add system instruction if provided
	if req.System != "" {
		geminiReq.SystemInstruction = &geminiContent{
			Parts: []geminiPart{{Text: req.System}},
		}
	}

	// Add generation config
	geminiReq.GenerationConfig = &geminiGenerationConfig{}
	if req.MaxTokens > 0 {
		geminiReq.GenerationConfig.MaxOutputTokens = req.MaxTokens
	} else {
		geminiReq.GenerationConfig.MaxOutputTokens = 8192
	}
	if req.Temperature > 0 {
		geminiReq.GenerationConfig.Temperature = req.Temperature
	}

	funcDecls := applyGeminiTools(&geminiReq, req.Tools, req.ToolChoice)
	if len(funcDecls) > 0 {
		log.Debug().Int("tool_count", len(funcDecls)).Strs("tool_names", func() []string {
			names := make([]string, len(funcDecls))
			for i, f := range funcDecls {
				names[i] = f.Name
			}
			return names
		}()).Msg("gemini request includes tools")
	}

	body, err := json.Marshal(geminiReq)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	// Build the URL with API key
	generateContentURL := fmt.Sprintf("%s/models/%s:generateContent?key=%s", c.baseURL, model, c.apiKey)

	log.Debug().Str("model", model).Str("base_url", c.baseURL).Msg("gemini Chat request")

	// Retry loop for transient errors
	var respBody []byte
	var lastErr error
	maxRetries := geminiMaxRetries

	for attempt := 0; attempt <= maxRetries; attempt++ {
		if attempt > 0 {
			// Exponential backoff: 2s, 4s, 8s
			backoff := geminiInitialBackoff * time.Duration(1<<(attempt-1))
			log.Warn().
				Int("attempt", attempt).
				Dur("backoff", backoff).
				Str("last_error", lastErr.Error()).
				Msg("Retrying Gemini API request after transient error")

			backoffTimer := time.NewTimer(backoff)
			select {
			case <-ctx.Done():
				if !backoffTimer.Stop() {
					select {
					case <-backoffTimer.C:
					default:
					}
				}
				return nil, ctx.Err()
			case <-backoffTimer.C:
			}
		}

		httpReq, err := http.NewRequestWithContext(ctx, "POST", generateContentURL, bytes.NewReader(body))
		if err != nil {
			return nil, fmt.Errorf("failed to create request: %w", err)
		}

		httpReq.Header.Set("Content-Type", "application/json")

		resp, err := c.client.Do(httpReq)
		if err != nil {
			// Check if this is a retryable connection error
			errStr := err.Error()
			if strings.Contains(errStr, "connection reset") ||
				strings.Contains(errStr, "connection refused") ||
				strings.Contains(errStr, "EOF") ||
				strings.Contains(errStr, "timeout") {
				lastErr = fmt.Errorf("connection error: %w", err)
				continue
			}
			return nil, fmt.Errorf("request failed: %w", err)
		}

		respBody, err = io.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			lastErr = fmt.Errorf("failed to read response: %w", err)
			continue
		}

		// Check for retryable HTTP errors
		if resp.StatusCode == 429 || resp.StatusCode == 503 || resp.StatusCode >= 500 {
			var errResp geminiError
			errMsg := string(respBody)
			if err := json.Unmarshal(respBody, &errResp); err == nil && errResp.Error.Message != "" {
				errMsg = errResp.Error.Message
			}
			errMsg = appendRateLimitInfo(errMsg, resp)
			lastErr = fmt.Errorf("API error (%d): %s", resp.StatusCode, errMsg)
			continue
		}

		// Non-retryable error
		if resp.StatusCode != http.StatusOK {
			var errResp geminiError
			if err := json.Unmarshal(respBody, &errResp); err == nil && errResp.Error.Message != "" {
				errMsg := appendRateLimitInfo(errResp.Error.Message, resp)
				return nil, fmt.Errorf("API error (%d): %s", resp.StatusCode, errMsg)
			}
			errMsg := appendRateLimitInfo(string(respBody), resp)
			return nil, fmt.Errorf("API error (%d): %s", resp.StatusCode, errMsg)
		}

		// Success - break out of retry loop
		lastErr = nil
		break
	}

	if lastErr != nil {
		return nil, fmt.Errorf("request failed after %d retries: %w", maxRetries, lastErr)
	}

	var geminiResp geminiResponse
	if err := json.Unmarshal(respBody, &geminiResp); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	// Check for prompt-level blocking
	if geminiResp.PromptFeedback != nil && geminiResp.PromptFeedback.BlockReason != "" {
		log.Warn().
			Str("block_reason", geminiResp.PromptFeedback.BlockReason).
			Msg("Gemini blocked the prompt")
		return nil, fmt.Errorf("prompt blocked by Gemini: %s", geminiResp.PromptFeedback.BlockReason)
	}

	if len(geminiResp.Candidates) == 0 {
		log.Warn().Str("raw_response", string(respBody)).Msg("gemini returned no candidates")
		return nil, fmt.Errorf("no response candidates returned")
	}

	candidate := geminiResp.Candidates[0]

	// Check for response-level blocking
	if candidate.FinishReason == "SAFETY" {
		blockedCategories := make([]string, 0)
		for _, safety := range candidate.SafetyRatings {
			if safety.Blocked {
				blockedCategories = append(blockedCategories, safety.Category)
			}
		}
		log.Warn().
			Strs("blocked_categories", blockedCategories).
			Msg("Gemini response blocked due to safety filters")
		return nil, fmt.Errorf("response blocked by Gemini safety filters: %v", blockedCategories)
	}

	// Extract content and tool calls from response
	var textContent string
	var toolCalls []ToolCall
	for _, part := range candidate.Content.Parts {
		if part.Text != "" {
			textContent += part.Text
		}
		if part.FunctionCall != nil {
			toolID := strings.TrimSpace(part.FunctionCall.ID)
			if toolID == "" {
				toolID = fmt.Sprintf("%s_%d", part.FunctionCall.Name, len(toolCalls))
			}
			signature := part.ThoughtSignature
			if len(signature) == 0 {
				signature = part.ThoughtSignatureSnake
			}
			toolCalls = append(toolCalls, ToolCall{
				ID:               toolID,
				Name:             part.FunctionCall.Name,
				Input:            part.FunctionCall.Args,
				ThoughtSignature: signature,
			})
		}
	}

	log.Debug().
		Str("model", model).
		Int("content_length", len(textContent)).
		Int("tool_calls", len(toolCalls)).
		Str("finish_reason", candidate.FinishReason).
		Msg("Gemini Chat response parsed")

	// Map finish reason - tool_use takes priority if there are tool calls
	stopReason := candidate.FinishReason
	if len(toolCalls) > 0 {
		// If there are tool calls, always signal tool_use so the agentic loop continues
		stopReason = "tool_use"
	} else if stopReason == "STOP" {
		stopReason = "end_turn"
	}

	var inputTokens, outputTokens int
	if geminiResp.UsageMetadata != nil {
		inputTokens = geminiResp.UsageMetadata.PromptTokenCount
		outputTokens = geminiResp.UsageMetadata.CandidatesTokenCount
	}

	return &ChatResponse{
		Content:      textContent,
		Model:        model,
		StopReason:   stopReason,
		ToolCalls:    toolCalls,
		InputTokens:  inputTokens,
		OutputTokens: outputTokens,
	}, nil
}

// TestConnection validates the API key by listing models
func (c *GeminiClient) TestConnection(ctx context.Context) error {
	if _, err := c.ListModels(ctx); err != nil {
		return fmt.Errorf("gemini test connection failed: %w", err)
	}
	return nil
}

// SupportsThinking returns true if the model supports extended thinking
func (c *GeminiClient) SupportsThinking(model string) bool {
	// Gemini models don't currently expose extended thinking in the streaming API
	return false
}

// geminiStreamEvent represents a streaming event from the Gemini API
type geminiStreamEvent struct {
	Candidates    []geminiCandidate    `json:"candidates,omitempty"`
	UsageMetadata *geminiUsageMetadata `json:"usageMetadata,omitempty"`
}

// ChatStream sends a chat request and streams the response via callback
func (c *GeminiClient) ChatStream(ctx context.Context, req ChatRequest, callback StreamCallback) error {
	// Convert messages to Gemini format (same as Chat)
	contents := convertMessagesToGemini(req.Messages)

	model := req.Model
	if strings.HasPrefix(model, "gemini:") {
		model = strings.TrimPrefix(model, "gemini:")
	}
	if model == "" {
		model = c.model
	}

	geminiReq := geminiRequest{
		Contents: contents,
	}

	if req.System != "" {
		geminiReq.SystemInstruction = &geminiContent{
			Parts: []geminiPart{{Text: req.System}},
		}
	}

	geminiReq.GenerationConfig = &geminiGenerationConfig{}
	if req.MaxTokens > 0 {
		geminiReq.GenerationConfig.MaxOutputTokens = req.MaxTokens
	} else {
		geminiReq.GenerationConfig.MaxOutputTokens = 8192
	}
	if req.Temperature > 0 {
		geminiReq.GenerationConfig.Temperature = req.Temperature
	}

	funcDecls := applyGeminiTools(&geminiReq, req.Tools, req.ToolChoice)
	if len(funcDecls) > 0 {
		toolNames := make([]string, len(funcDecls))
		for i, f := range funcDecls {
			toolNames[i] = f.Name
		}
		log.Debug().
			Int("tool_count", len(funcDecls)).
			Strs("tool_names", toolNames).
			Msg("Gemini stream request includes tools")
	}

	body, err := json.Marshal(geminiReq)
	if err != nil {
		return fmt.Errorf("failed to marshal request: %w", err)
	}

	// Log the full request body for debugging (at trace level to avoid noise)
	log.Trace().
		Str("model", model).
		RawJSON("request_body", body).
		Msg("Gemini stream request body")

	// Use streamGenerateContent endpoint for streaming
	streamGenerateContentURL := fmt.Sprintf("%s/models/%s:streamGenerateContent?key=%s&alt=sse", c.baseURL, model, c.apiKey)

	// Retry loop for transient errors (matching non-streaming Chat behavior)
	var resp *http.Response
	var lastErr error
	maxRetries := geminiMaxRetries

	for attempt := 0; attempt <= maxRetries; attempt++ {
		// Bail out early if the parent context is already cancelled
		if ctx.Err() != nil {
			if lastErr != nil {
				return fmt.Errorf("request failed (context cancelled after %d attempts): %w", attempt, lastErr)
			}
			return ctx.Err()
		}

		if attempt > 0 {
			backoff := geminiInitialBackoff * time.Duration(1<<(attempt-1))
			log.Warn().
				Int("attempt", attempt).
				Dur("backoff", backoff).
				Str("last_error", lastErr.Error()).
				Msg("Retrying Gemini stream request after transient error")

			backoffTimer := time.NewTimer(backoff)
			select {
			case <-ctx.Done():
				if !backoffTimer.Stop() {
					select {
					case <-backoffTimer.C:
					default:
					}
				}
				return ctx.Err()
			case <-backoffTimer.C:
			}
		}

		httpReq, err := http.NewRequestWithContext(ctx, "POST", streamGenerateContentURL, bytes.NewReader(body))
		if err != nil {
			return fmt.Errorf("failed to create request: %w", err)
		}

		httpReq.Header.Set("Content-Type", "application/json")
		httpReq.Header.Set("Accept", "text/event-stream")

		resp, err = c.client.Do(httpReq)
		if err != nil {
			errStr := err.Error()
			if strings.Contains(errStr, "connection reset") ||
				strings.Contains(errStr, "connection refused") ||
				strings.Contains(errStr, "EOF") ||
				strings.Contains(errStr, "timeout") ||
				strings.Contains(errStr, "deadline exceeded") {
				lastErr = fmt.Errorf("connection error: %w", err)
				continue
			}
			return fmt.Errorf("request failed: %w", err)
		}

		if resp.StatusCode == 429 || resp.StatusCode == 503 || resp.StatusCode >= 500 {
			respBody, _ := io.ReadAll(resp.Body)
			resp.Body.Close()
			var errResp geminiError
			errMsg := string(respBody)
			if err := json.Unmarshal(respBody, &errResp); err == nil && errResp.Error.Message != "" {
				errMsg = errResp.Error.Message
			}
			errMsg = appendRateLimitInfo(errMsg, resp)
			lastErr = fmt.Errorf("API error (%d): %s", resp.StatusCode, errMsg)
			continue
		}

		lastErr = nil
		break
	}

	if lastErr != nil {
		return fmt.Errorf("request failed after %d retries: %w", maxRetries, lastErr)
	}

	if resp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		var errResp geminiError
		if err := json.Unmarshal(respBody, &errResp); err == nil && errResp.Error.Message != "" {
			errMsg := appendRateLimitInfo(errResp.Error.Message, resp)
			return fmt.Errorf("API error (%d): %s", resp.StatusCode, errMsg)
		}
		errMsg := appendRateLimitInfo(string(respBody), resp)
		return fmt.Errorf("API error (%d): %s", resp.StatusCode, errMsg)
	}
	defer resp.Body.Close()

	// Parse SSE stream
	reader := resp.Body
	buf := make([]byte, 4096)
	var pendingData string
	var toolCalls []ToolCall
	var inputTokens, outputTokens int
	var finishReason string

	for {
		n, err := reader.Read(buf)
		if n > 0 {
			pendingData += string(buf[:n])
			lines := strings.Split(pendingData, "\n")

			// Keep the last incomplete line for next iteration
			pendingData = lines[len(lines)-1]
			lines = lines[:len(lines)-1]

			for _, line := range lines {
				line = strings.TrimSpace(line)

				if !strings.HasPrefix(line, "data:") {
					continue
				}

				data := strings.TrimPrefix(line, "data:")
				data = strings.TrimSpace(data)

				if data == "" {
					continue
				}

				var event geminiStreamEvent
				if err := json.Unmarshal([]byte(data), &event); err != nil {
					log.Debug().Err(err).Str("data", data).Msg("failed to parse Gemini stream event")
					continue
				}

				if event.UsageMetadata != nil {
					inputTokens = event.UsageMetadata.PromptTokenCount
					outputTokens = event.UsageMetadata.CandidatesTokenCount
				}

				for _, candidate := range event.Candidates {
					if candidate.FinishReason != "" {
						finishReason = candidate.FinishReason
					}

					for _, part := range candidate.Content.Parts {
						if part.Text != "" {
							callback(StreamEvent{
								Type: "content",
								Data: ContentEvent{Text: part.Text},
							})
						}

						if part.FunctionCall != nil {
							toolID := strings.TrimSpace(part.FunctionCall.ID)
							if toolID == "" {
								toolID = fmt.Sprintf("%s_%d", part.FunctionCall.Name, len(toolCalls))
							}
							signature := part.ThoughtSignature
							if len(signature) == 0 {
								signature = part.ThoughtSignatureSnake
							}
							log.Debug().
								Str("tool_name", part.FunctionCall.Name).
								Interface("tool_args", part.FunctionCall.Args).
								Msg("Gemini called tool")
							callback(StreamEvent{
								Type: "tool_start",
								Data: ToolStartEvent{
									ID:    toolID,
									Name:  part.FunctionCall.Name,
									Input: part.FunctionCall.Args,
								},
							})
							toolCalls = append(toolCalls, ToolCall{
								ID:               toolID,
								Name:             part.FunctionCall.Name,
								Input:            part.FunctionCall.Args,
								ThoughtSignature: signature,
							})
						}
					}
				}
			}
		}
		if err != nil {
			if err == io.EOF {
				break
			}
			return fmt.Errorf("stream read error: %w", err)
		}
	}

	// Send done event
	stopReason := finishReason
	if len(toolCalls) > 0 {
		stopReason = "tool_use"
	} else if stopReason == "STOP" {
		stopReason = "end_turn"
	}

	callback(StreamEvent{
		Type: "done",
		Data: DoneEvent{
			StopReason:   stopReason,
			ToolCalls:    toolCalls,
			InputTokens:  inputTokens,
			OutputTokens: outputTokens,
		},
	})

	return nil
}

// ListModels fetches available models from the Gemini API
func (c *GeminiClient) ListModels(ctx context.Context) ([]ModelInfo, error) {
	modelsURL := fmt.Sprintf("%s/models?key=%s", c.baseURL, c.apiKey)

	req, err := http.NewRequestWithContext(ctx, "GET", modelsURL, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		errMsg := appendRateLimitInfo(string(body), resp)
		return nil, fmt.Errorf("API error (%d): %s", resp.StatusCode, errMsg)
	}

	var result struct {
		Models []struct {
			Name                       string   `json:"name"`
			DisplayName                string   `json:"displayName"`
			Description                string   `json:"description"`
			SupportedGenerationMethods []string `json:"supportedGenerationMethods"`
		} `json:"models"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	models := make([]ModelInfo, 0, len(result.Models))
	cache := GetNotableCache()
	for _, m := range result.Models {
		// Only include models that support generateContent (chat)
		supportsChat := false
		for _, method := range m.SupportedGenerationMethods {
			if method == "generateContent" {
				supportsChat = true
				break
			}
		}
		if !supportsChat {
			continue
		}

		// Extract model ID from the full name (e.g., "models/gemini-1.5-pro" -> "gemini-1.5-pro")
		modelID := strings.TrimPrefix(m.Name, "models/")

		models = append(models, ModelInfo{
			ID:          modelID,
			Name:        m.DisplayName,
			Description: m.Description,
			Notable:     cache.IsNotable("gemini", modelID, 0),
		})
	}

	return models, nil
}