package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"

	"cloud.google.com/go/vertexai/genai"
	"google.golang.org/api/option"
	htransport "google.golang.org/api/transport/http"
)

func main() {
	ctx := context.Background()
	// Make a transport that has all the necessary authentication and other features for Google Cloud.
	transport, err := htransport.NewTransport(ctx, http.DefaultTransport,
		option.WithScopes("https://www.googleapis.com/auth/cloud-platform"))
	if err != nil {
		log.Fatal(err)
	}
	// Wrap it in a headerTransport (see definition below), which can be used to add any headers
	// to a request. (It will also overwrite existing headers, so use with caution.)
	headers := http.Header{}
	headers.Set("My-Header", "17")
	transport = &headerTransport{
		headers: headers,
		base:    transport,
	}

	// Now make an http.Client with our transport and pass it as an option.
	hc := &http.Client{Transport: transport}
	client, err := genai.NewClient(ctx, os.Getenv("PROJECT"), "us-central1", genai.WithREST(), option.WithHTTPClient(hc))
	if err != nil {
		log.Fatal(err)
	}
	defer client.Close()
	model := client.GenerativeModel("gemini-pro")
	res, err := model.GenerateContent(ctx, genai.Text("My prompt."))
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println(res.Candidates[0].Content.Parts)
}

type headerTransport struct {
	headers http.Header
	base    http.RoundTripper
}

func (t *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	rt := t.base
	newReq := *req
	newReq.Header = make(http.Header)
	for k, vv := range req.Header {
		newReq.Header[k] = vv
	}

	for k, v := range t.headers {
		newReq.Header[k] = v
	}

	return rt.RoundTrip(&newReq)
}