package gemini

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

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

// Client defines the common interface for Gemini interactions
type Client interface {
	Generate(ctx context.Context, prompt string) (string, error)
	Close() error
}

// --- REST Client (for custom BaseURL / proxy like lar) ---

type RestClient struct {
	httpClient *http.Client
	baseURL    string
	model      string
}

func NewRestClient(apiKey, baseURL, model string) (*RestClient, error) {
	baseURL = strings.TrimSuffix(baseURL, "/")
	return &RestClient{
		httpClient: &http.Client{},
		baseURL:    baseURL,
		model:      model,
	}, nil
}

type restRequest struct {
	Contents []restContent `json:"contents"`
}

type restContent struct {
	Parts []restPart `json:"parts"`
}

type restPart struct {
	Text string `json:"text"`
}

type restResponse struct {
	Candidates []restCandidate `json:"candidates"`
}

type restCandidate struct {
	Content restContent `json:"content"`
}

func (c *RestClient) Generate(ctx context.Context, prompt string) (string, error) {
	url := fmt.Sprintf("%s/v1beta/models/%s:generateContent", c.baseURL, c.model)

	reqBody := restRequest{
		Contents: []restContent{
			{
				Parts: []restPart{
					{Text: prompt},
				},
			},
		},
	}

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

	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonBytes))
	if err != nil {
		return "", fmt.Errorf("failed to create http request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	
	// Note: Authentication (x-goog-api-key) is injected transparently by the lar proxy transport.

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

	bodyBytes, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("gemini rest api error (status %d): %s", resp.StatusCode, string(bodyBytes))
	}

	var res restResponse
	if err := json.Unmarshal(bodyBytes, &res); err != nil {
		return "", fmt.Errorf("failed to decode rest response: %w (body: %s)", err, string(bodyBytes))
	}

	if len(res.Candidates) == 0 || len(res.Candidates[0].Content.Parts) == 0 {
		return "", fmt.Errorf("empty candidates in rest response")
	}

	return res.Candidates[0].Content.Parts[0].Text, nil
}

func (c *RestClient) Close() error {
	return nil
}


// --- gRPC Client (official SDK for direct Google API connection) ---

type GrpcClient struct {
	client *genai.Client
	model  string
}

func NewGrpcClient(ctx context.Context, apiKey, model string) (*GrpcClient, error) {
	c, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
	if err != nil {
		return nil, err
	}
	return &GrpcClient{
		client: c,
		model:  model,
	}, nil
}

func (c *GrpcClient) Generate(ctx context.Context, prompt string) (string, error) {
	model := c.client.GenerativeModel(c.model)
	resp, err := model.GenerateContent(ctx, genai.Text(prompt))
	if err != nil {
		return "", err
	}

	if len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 {
		return "", fmt.Errorf("empty response from Gemini gRPC")
	}

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

	return fmt.Sprintf("%v", part), nil
}

func (c *GrpcClient) Close() error {
	return c.client.Close()
}


// --- Factory helper ---

func NewClient(ctx context.Context, apiKey, baseURL, model string) (Client, error) {
	if baseURL != "" {
		return NewRestClient(apiKey, baseURL, model)
	}
	return NewGrpcClient(ctx, apiKey, model)
}