package service

import (
	"context"
	"sync"

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

type GeminiService struct {
	client *genai.Client
	model  *genai.GenerativeModel
	apiKey string
	mu     sync.RWMutex
}

func NewGeminiService(apiKey string) (*GeminiService, error) {
	ctx := context.Background()
	client, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
	if err != nil {
		return nil, errors.Wrap(err, "failed to create Gemini client")
	}

	return &GeminiService{
		client: client,
		model:  client.GenerativeModel("gemini-pro"),
		apiKey: apiKey,
	}, nil
}

func (s *GeminiService) ProcessText(ctx context.Context, text string) (string, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	resp, err := s.model.GenerateContent(ctx, genai.Text(text))
	if err != nil {
		return "", errors.Wrap(err, "failed to generate content")
	}

	if len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 {
		return "", errors.New("no response from Gemini")
	}

	if textPart, ok := resp.Candidates[0].Content.Parts[0].(genai.Text); ok {
		return string(textPart), nil
	}

	return "", errors.New("invalid response format from Gemini")
}

// Close implements graceful shutdown
func (s *GeminiService) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.client != nil {
		s.client.Close()
	}
	return nil
}