// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tools

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"testing"

	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
	"google.golang.org/genai"
)

type mockModels struct{}

// Mock Vertex AI Search (VAIS) response
func (m *mockModels) GenerateContentMock(ctx context.Context, model string, contents []*genai.Content, cfg *genai.GenerateContentConfig) (*genai.GenerateContentResponse, error) {
	return &genai.GenerateContentResponse{
		Candidates: []*genai.Candidate{
			{
				Content: &genai.Content{
					Parts: []*genai.Part{{Text: "Mocked VAIS search result: Found relevant content for your query."}},
				},
			},
		},
	}, nil
}

func generateWithGoogleVAISMock(w io.Writer, dataStore string) error {
	models := &mockModels{}
	ctx := context.Background()

	resp, err := models.GenerateContentMock(ctx,
		"google/gemini-2.0-flash",
		[]*genai.Content{
			{
				Role: "user",
				Parts: []*genai.Part{
					{Text: fmt.Sprintf("Search for insights using datastore: %s", dataStore)},
				},
			},
		},
		nil,
	)
	if err != nil {
		return fmt.Errorf("failed to generate content: %w", err)
	}

	fmt.Fprintln(w, resp.Text())
	return nil
}

func TestTextGeneration(t *testing.T) {
	tc := testutil.SystemTest(t)

	t.Setenv("GOOGLE_GENAI_USE_VERTEXAI", "1")
	t.Setenv("GOOGLE_CLOUD_LOCATION", "us-central1")
	t.Setenv("GOOGLE_CLOUD_PROJECT", tc.ProjectID)

	buf := new(bytes.Buffer)

	t.Run("generate with code execution tool", func(t *testing.T) {
		buf.Reset()
		err := generateWithCodeExec(buf)
		if err != nil {
			t.Fatalf("generateWithCodeExec failed: %v", err)
		}

		output := buf.String()
		if output == "" {
			t.Error("expected non-empty output, got empty")
		}
	})

	t.Run("generate with func declaration and func response", func(t *testing.T) {
		buf.Reset()
		err := generateWithFuncCall(buf)
		if err != nil {
			t.Fatalf("generateWithFuncCall failed: %v", err)
		}

		output := buf.String()
		if output == "" {
			t.Error("expected non-empty output, got empty")
		}
	})

	t.Run("generate with Google Search", func(t *testing.T) {
		buf.Reset()
		err := generateWithGoogleSearch(buf)
		if err != nil {
			t.Fatalf("generateWithGoogleSearch failed: %v", err)
		}

		output := buf.String()
		if output == "" {
			t.Error("expected non-empty output, got empty")
		}
	})

	t.Run("generate with VAIS Search", func(t *testing.T) {
		buf.Reset()
		dataStore := fmt.Sprintf("projects/%s/locations/global/collections/default_collection/dataStores/grounding-test-datastore", tc.ProjectID)

		err := generateWithGoogleVAISMock(buf, dataStore)
		if err != nil {
			t.Fatalf("generateWithGoogleVAIS failed: %v", err)
		}

		output := buf.String()
		if output == "" {
			t.Error("expected non-empty output, got empty")
		}
	})

	t.Run("generate content with maps coordinates", func(t *testing.T) {
		buf.Reset()
		if err := generateGoogleMapsCoordinatesWithText(buf); err != nil {
			t.Fatalf("generateGmapsCoordinatesWithText failed: %v", err)
		}

		output := buf.String()
		if output == "" {
			t.Error("expected non-empty output, got empty")
		}
	})

	t.Run("generate content with url context", func(t *testing.T) {
		buf.Reset()
		if err := generateURLContentWithText(buf); err != nil {
			t.Fatalf("generateURLContentWithText failed: %v", err)
		}

		output := buf.String()
		if output == "" {
			t.Error("expected non-empty output, got empty")
		}
	})

	t.Run("generate content with google search and url context", func(t *testing.T) {
		buf.Reset()
		if err := generateGoogleSearchURLContextWithText(buf); err != nil {
			t.Fatalf("generateGSearchURLContentWithText failed: %v", err)
		}

		output := buf.String()
		if output == "" {
			t.Error("expected non-empty output, got empty")
		}
	})

	t.Run("generate with Local img and code exec", func(t *testing.T) {
		buf.Reset()
		err := generateWithLocalImgAndCodeExec(buf)
		if err != nil {
			t.Fatalf("generateWithLocalImgAndCodeExec failed: %v", err)
		}

		output := buf.String()
		if output == "" {
			t.Error("expected non-empty output, got empty")
		}
	})
}
