package llm_cli

import (
	"context"
	"time"

	"github.com/tidbcloud/devops/pkg/infra/utils/pointer"
	"google.golang.org/genai"
)

func (llm *LLMCli) SimpleSearch(input string) (string, error) {
	// Set 2 minute timeout for LLM search with GoogleSearch tool
	// Web searches can take longer due to webpage fetching and verification
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()

	result, err := llm.Gemini.Models.GenerateContent(
		ctx,
		"gemini-flash-latest",
		genai.Text(input),
		&genai.GenerateContentConfig{
			Temperature: pointer.Pointer(float32(0.0)),
			Tools: []*genai.Tool{
				{
					GoogleSearch: &genai.GoogleSearch{},
				},
			},
		},
	)
	if err != nil {
		return "", err
	}
	return result.Text(), nil
}

func (llm *LLMCli) GenerateText(ctx context.Context, input string) (string, error) {
	result, err := llm.Gemini.Models.GenerateContent(
		ctx,
		"gemini-flash-latest",
		genai.Text(input),
		&genai.GenerateContentConfig{
			Temperature: pointer.Pointer(float32(0.0)),
		},
	)
	if err != nil {
		return "", err
	}
	return result.Text(), nil
}

func NewGeminiClient(apiKey string) (*genai.Client, error) {
	return genai.NewClient(
		context.Background(),
		&genai.ClientConfig{
			Backend: genai.BackendGeminiAPI,
			APIKey:  apiKey,
		},
	)
}