package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// provider_gemini.go is the Gemini backend. Two faces:
//
//   - geminiClient: a thin LLMClient over the OpenAI-compatible endpoint, used for the
//     cheap trivial route and the Layer-1 router classifier. Same wire format as Grok,
//     so it reuses the shared transport (httpllm.go).
//   - groundedSearch: a SEPARATE call against the NATIVE v1beta generateContent endpoint
//     with the google_search tool. Grounding does NOT work on the OpenAI-compat layer —
//     it is silently ignored THERE (F-EXT-3, an endpoint limitation, NOT a model-version
//     one: the google_search tool is supported by current models including
//     gemini-2.5-flash-lite per ai.google.dev). So the web layer that wants Gemini
//     grounding must use this native path and VERIFY citations came back, else degrade.

// maxGroundingRespBytes caps the native grounding response read so a hostile or
// malfunctioning endpoint cannot drive unbounded memory growth via io.ReadAll.
const maxGroundingRespBytes = 1 << 20 // 1 MiB

type geminiClient struct {
	http       *openAIClient
	nativeBase string // …/v1beta — derived from the OpenAI-compat base by dropping /openai
	key        string
	model      string
	httpc      *http.Client
	log        *slog.Logger
}

// NewGeminiClient builds the Gemini backend. base is the OpenAI-compatible endpoint
// (…/v1beta/openai); the native grounding endpoint is derived from it. Returns the
// concrete type (not just LLMClient) because the web layer needs groundedSearch too.
func NewGeminiClient(base, key, model string, logger *slog.Logger) *geminiClient {
	return &geminiClient{
		http:       newOpenAIClient("gemini", base, key, nil, logger),
		nativeBase: strings.TrimSuffix(base, "/openai"),
		key:        key,
		model:      model,
		httpc: &http.Client{
			// Refuse redirects: the native endpoint never legitimately redirects,
			// and following one to another host would be an exfil path.
			CheckRedirect: func(*http.Request, []*http.Request) error {
				return http.ErrUseLastResponse
			},
		},
		log: logger,
	}
}

// Complete answers via the OpenAI-compatible endpoint (trivial route + classifier).
func (c *geminiClient) Complete(ctx context.Context, req LLMRequest) (*LLMResponse, error) {
	msgs := make([]openAIMessage, len(req.Messages))
	for i, m := range req.Messages {
		msgs[i] = openAIMessage{Role: m.Role, Content: m.Content}
	}
	var respFormat any
	if req.JSONOnly {
		respFormat = map[string]string{"type": "json_object"}
	}
	resp, err := c.http.complete(ctx, openAIRequest{
		Model:          req.Model,
		Messages:       msgs,
		MaxTokens:      req.MaxTokens,
		Temperature:    req.Temperature,
		Stream:         false,
		ResponseFormat: respFormat,
	}, nil)
	if err != nil {
		return nil, err
	}
	return &LLMResponse{
		Text: resp.Text(),
		Usage: Usage{
			PromptTokens:     resp.Usage.PromptTokens,
			CachedTokens:     resp.Usage.PromptTokensDetails.CachedTokens,
			CompletionTokens: resp.Usage.CompletionTokens,
		},
		ProviderRequestID: resp.ID,
	}, nil
}

// --- native v1beta grounded search (google_search tool) ---------------------------

type geminiGroundResult struct {
	Digest    string
	Citations []string    // redirect URIs — the verify-gate + citation_count
	Sources   []WebSource // the same chunks with their publisher-domain titles (web.title)
	Usage     Usage
}

// native generateContent wire types (only the fields we read/write).
type geminiNativeRequest struct {
	Contents         []geminiContent  `json:"contents"`
	Tools            []geminiTool     `json:"tools"`
	GenerationConfig *geminiGenConfig `json:"generationConfig,omitempty"`
}
type geminiGenConfig struct {
	MaxOutputTokens int `json:"maxOutputTokens,omitempty"`
}
type geminiContent struct {
	Role  string       `json:"role,omitempty"`
	Parts []geminiPart `json:"parts"`
}
type geminiPart struct {
	Text string `json:"text"`
}
type geminiTool struct {
	// google_search is the current grounding tool (all current models, incl. the 2.5
	// family; legacy models used google_search_retrieval). The empty object enables it.
	GoogleSearch struct{} `json:"google_search"`
}
type geminiNativeResponse struct {
	Candidates []struct {
		Content struct {
			Parts []geminiPart `json:"parts"`
		} `json:"content"`
		GroundingMetadata struct {
			GroundingChunks []struct {
				Web struct {
					URI   string `json:"uri"`
					Title string `json:"title"`
				} `json:"web"`
			} `json:"groundingChunks"`
		} `json:"groundingMetadata"`
	} `json:"candidates"`
	UsageMetadata struct {
		PromptTokenCount        int `json:"promptTokenCount"`
		CandidatesTokenCount    int `json:"candidatesTokenCount"`
		CachedContentTokenCount int `json:"cachedContentTokenCount"`
		// ThoughtsTokenCount is ADDITIVE to candidatesTokenCount (Google documents
		// thinking tokens separately) — same convention as Usage.ReasoningTokens.
		// 0 on flash-lite (thinking off by default), but a model swap must not
		// silently re-open the unbilled-thinking hole.
		ThoughtsTokenCount int `json:"thoughtsTokenCount"`
	} `json:"usageMetadata"`
}

// groundedDigestMaxTokens bounds the fetched digest. Big enough for a dense
// multi-fact summary, small enough that the synthesis prompt stays cheap.
const groundedDigestMaxTokens = 1024

// groundedSearch runs one grounded generateContent against the native endpoint and
// returns the model's grounded answer plus the source URLs. It REQUIRES citations:
// if groundingMetadata has no chunks the request was not actually grounded (the
// silent-ignore failure mode, F-EXT-3), so it errors and the caller degrades rather
// than passing off ungrounded — possibly stale — text as fresh.
//
// The query is wrapped in a fetch instruction instead of being sent bare: a bare
// query leaves digest depth, dating, and language to Gemini's chat defaults, and the
// model has no idea what "today" is. The digest is the route's most important
// intermediate artifact and its tokens are ~free at flash-lite prices, so instruct it.
func (c *geminiClient) groundedSearch(ctx context.Context, query string) (geminiGroundResult, error) {
	prompt := dateNote(time.Now()) + webFetchInstruction + query
	body, err := json.Marshal(geminiNativeRequest{
		Contents:         []geminiContent{{Role: "user", Parts: []geminiPart{{Text: prompt}}}},
		Tools:            []geminiTool{{}},
		GenerationConfig: &geminiGenConfig{MaxOutputTokens: groundedDigestMaxTokens},
	})
	if err != nil {
		return geminiGroundResult{}, err
	}

	// The key goes in the x-goog-api-key header, NOT the query string: a
	// transport-level failure (DNS/timeout/TLS) returns a *url.Error whose
	// Error() embeds the request URL verbatim, and the cascade logs that error
	// at WARN — a `?key=` would leak the secret into the log backend.
	endpoint := fmt.Sprintf("%s/models/%s:generateContent",
		c.nativeBase, url.PathEscape(c.model))

	reqCtx, cancel := context.WithTimeout(ctx, 15*time.Second) // web/grounding budget (§8.2.2)
	defer cancel()
	req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, endpoint, bytes.NewReader(body))
	if err != nil {
		return geminiGroundResult{}, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-goog-api-key", c.key)

	resp, err := c.httpc.Do(req)
	if err != nil {
		return geminiGroundResult{}, err
	}
	defer resp.Body.Close()
	data, _ := io.ReadAll(io.LimitReader(resp.Body, maxGroundingRespBytes))
	logLLMExchange(ctx, c.log, "gemini_grounding", body, resp.StatusCode, data)
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return geminiGroundResult{}, fmt.Errorf("gemini grounding http %d: %s", resp.StatusCode, snippet(data))
	}

	var out geminiNativeResponse
	if err := json.Unmarshal(data, &out); err != nil {
		return geminiGroundResult{}, fmt.Errorf("gemini grounding decode: %w", err)
	}
	if len(out.Candidates) == 0 {
		return geminiGroundResult{}, fmt.Errorf("gemini grounding: no candidates")
	}

	var sb strings.Builder
	for _, p := range out.Candidates[0].Content.Parts {
		sb.WriteString(p.Text)
	}
	var citations []string
	var sources []WebSource
	for _, ch := range out.Candidates[0].GroundingMetadata.GroundingChunks {
		if ch.Web.URI != "" {
			citations = append(citations, ch.Web.URI)
			// web.uri is the grounding-api-redirect (NOT the publisher URL — and Gemini's
			// terms forbid resolving it server-side); web.title is the publisher domain
			// ("rbc.ru"). Keep both: the user clicks the redirect to reach the real article.
			sources = append(sources, WebSource{Title: ch.Web.Title, URL: ch.Web.URI})
		}
	}
	// The verify-gate: no citations ⇒ not actually grounded ⇒ degrade.
	if len(citations) == 0 {
		return geminiGroundResult{}, fmt.Errorf("gemini grounding: no citations (ungrounded — degrade)")
	}
	return geminiGroundResult{
		Digest:    strings.TrimSpace(sb.String()),
		Citations: citations,
		Sources:   sources,
		Usage: Usage{
			PromptTokens:     out.UsageMetadata.PromptTokenCount,
			CachedTokens:     out.UsageMetadata.CachedContentTokenCount,
			CompletionTokens: out.UsageMetadata.CandidatesTokenCount,
			ReasoningTokens:  out.UsageMetadata.ThoughtsTokenCount,
		},
	}, nil
}