// Copyright 2024 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 neo4j

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"os"
	"reflect"
	"regexp"
	"strings"
	"testing"
	"time"

	"github.com/neo4j/neo4j-go-driver/v5/neo4j"

	"github.com/googleapis/genai-toolbox/internal/testutils"
	"github.com/googleapis/genai-toolbox/tests"
)

var (
	Neo4jSourceKind = "neo4j"
	Neo4jDatabase   = os.Getenv("NEO4J_DATABASE")
	Neo4jUri        = os.Getenv("NEO4J_URI")
	Neo4jUser       = os.Getenv("NEO4J_USER")
	Neo4jPass       = os.Getenv("NEO4J_PASS")
)

// getNeo4jVars retrieves necessary Neo4j connection details from environment variables.
// It fails the test if any required variable is not set.
func getNeo4jVars(t *testing.T) map[string]any {
	switch "" {
	case Neo4jDatabase:
		t.Fatal("'NEO4J_DATABASE' not set")
	case Neo4jUri:
		t.Fatal("'NEO4J_URI' not set")
	case Neo4jUser:
		t.Fatal("'NEO4J_USER' not set")
	case Neo4jPass:
		t.Fatal("'NEO4J_PASS' not set")
	}

	return map[string]any{
		"kind":     Neo4jSourceKind,
		"uri":      Neo4jUri,
		"database": Neo4jDatabase,
		"user":     Neo4jUser,
		"password": Neo4jPass,
	}
}

// TestNeo4jToolEndpoints sets up an integration test server and tests the API endpoints
// for various Neo4j tools, including cypher execution and schema retrieval.
func TestNeo4jToolEndpoints(t *testing.T) {
	sourceConfig := getNeo4jVars(t)
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()

	var args []string

	// Write config into a file and pass it to the command.
	// This configuration defines the data source and the tools to be tested.
	toolsFile := map[string]any{
		"sources": map[string]any{
			"my-neo4j-instance": sourceConfig,
		},
		"tools": map[string]any{
			"my-simple-cypher-tool": map[string]any{
				"kind":        "neo4j-cypher",
				"source":      "my-neo4j-instance",
				"description": "Simple tool to test end to end functionality.",
				"statement":   "RETURN 1 as a;",
			},
			"my-simple-execute-cypher-tool": map[string]any{
				"kind":        "neo4j-execute-cypher",
				"source":      "my-neo4j-instance",
				"description": "Simple tool to test end to end functionality.",
			},
			"my-readonly-execute-cypher-tool": map[string]any{
				"kind":        "neo4j-execute-cypher",
				"source":      "my-neo4j-instance",
				"description": "A readonly cypher execution tool.",
				"readOnly":    true,
			},
			"my-schema-tool": map[string]any{
				"kind":        "neo4j-schema",
				"source":      "my-neo4j-instance",
				"description": "A tool to get the Neo4j schema.",
			},
			"my-schema-tool-with-cache": map[string]any{
				"kind":               "neo4j-schema",
				"source":             "my-neo4j-instance",
				"description":        "A schema tool with a custom cache expiration.",
				"cacheExpireMinutes": 10,
			},
			"my-populated-schema-tool": map[string]any{
				"kind":        "neo4j-schema",
				"source":      "my-neo4j-instance",
				"description": "A tool to get the Neo4j schema from a populated DB.",
			},
		},
	}
	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
	if err != nil {
		t.Fatalf("command initialization returned an error: %s", err)
	}
	defer cleanup()

	waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
	if err != nil {
		t.Logf("toolbox command logs: \n%s", out)
		t.Fatalf("toolbox didn't start successfully: %s", err)
	}

	// Test tool `GET` endpoints to verify their manifests are correct.
	tcs := []struct {
		name string
		api  string
		want map[string]any
	}{
		{
			name: "get my-simple-cypher-tool",
			api:  "http://127.0.0.1:5000/api/tool/my-simple-cypher-tool/",
			want: map[string]any{
				"my-simple-cypher-tool": map[string]any{
					"description":  "Simple tool to test end to end functionality.",
					"parameters":   []any{},
					"authRequired": []any{},
				},
			},
		},
		{
			name: "get my-simple-execute-cypher-tool",
			api:  "http://127.0.0.1:5000/api/tool/my-simple-execute-cypher-tool/",
			want: map[string]any{
				"my-simple-execute-cypher-tool": map[string]any{
					"description": "Simple tool to test end to end functionality.",
					"parameters": []any{
						map[string]any{
							"name":        "cypher",
							"type":        "string",
							"required":    true,
							"description": "The cypher to execute.",
							"authSources": []any{},
						},
					},
					"authRequired": []any{},
				},
			},
		},
		{
			name: "get my-schema-tool",
			api:  "http://127.0.0.1:5000/api/tool/my-schema-tool/",
			want: map[string]any{
				"my-schema-tool": map[string]any{
					"description":  "A tool to get the Neo4j schema.",
					"parameters":   []any{},
					"authRequired": []any{},
				},
			},
		},
		{
			name: "get my-schema-tool-with-cache",
			api:  "http://127.0.0.1:5000/api/tool/my-schema-tool-with-cache/",
			want: map[string]any{
				"my-schema-tool-with-cache": map[string]any{
					"description":  "A schema tool with a custom cache expiration.",
					"parameters":   []any{},
					"authRequired": []any{},
				},
			},
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			resp, err := http.Get(tc.api)
			if err != nil {
				t.Fatalf("error when sending a request: %s", err)
			}
			defer resp.Body.Close()
			if resp.StatusCode != 200 {
				t.Fatalf("response status code is not 200")
			}

			var body map[string]interface{}
			err = json.NewDecoder(resp.Body).Decode(&body)
			if err != nil {
				t.Fatalf("error parsing response body")
			}

			got, ok := body["tools"]
			if !ok {
				t.Fatalf("unable to find tools in response body")
			}
			if !reflect.DeepEqual(got, tc.want) {
				t.Fatalf("got %q, want %q", got, tc.want)
			}
		})
	}

	// Test tool `invoke` endpoints to verify their functionality.
	invokeTcs := []struct {
		name               string
		api                string
		requestBody        io.Reader
		want               string
		wantStatus         int
		wantErrorSubstring string
		prepareData        func(t *testing.T)
		validateFunc       func(t *testing.T, body string)
	}{
		{
			name:        "invoke my-simple-cypher-tool",
			api:         "http://127.0.0.1:5000/api/tool/my-simple-cypher-tool/invoke",
			requestBody: bytes.NewBuffer([]byte(`{}`)),
			want:        "[{\"a\":1}]",
			wantStatus:  http.StatusOK,
		},
		{
			name:        "invoke my-simple-execute-cypher-tool",
			api:         "http://127.0.0.1:5000/api/tool/my-simple-execute-cypher-tool/invoke",
			requestBody: bytes.NewBuffer([]byte(`{"cypher": "RETURN 1 as a;"}`)),
			want:        "[{\"a\":1}]",
			wantStatus:  http.StatusOK,
		},
		{
			name:               "invoke readonly tool with write query",
			api:                "http://127.0.0.1:5000/api/tool/my-readonly-execute-cypher-tool/invoke",
			requestBody:        bytes.NewBuffer([]byte(`{"cypher": "CREATE (n:TestNode)"}`)),
			wantStatus:         http.StatusBadRequest,
			wantErrorSubstring: "this tool is read-only and cannot execute write queries",
		},
		{
			name:        "invoke my-schema-tool",
			api:         "http://127.0.0.1:5000/api/tool/my-schema-tool/invoke",
			requestBody: bytes.NewBuffer([]byte(`{}`)),
			wantStatus:  http.StatusOK,
			validateFunc: func(t *testing.T, body string) {
				var result map[string]any
				if err := json.Unmarshal([]byte(body), &result); err != nil {
					t.Fatalf("failed to unmarshal schema result: %v", err)
				}
				// Check for the presence of top-level keys in the schema response.
				expectedKeys := []string{"nodeLabels", "relationships", "constraints", "indexes", "databaseInfo", "statistics"}
				for _, key := range expectedKeys {
					if _, ok := result[key]; !ok {
						t.Errorf("expected key %q not found in schema response", key)
					}
				}
			},
		},
		{
			name:        "invoke my-schema-tool-with-cache",
			api:         "http://127.0.0.1:5000/api/tool/my-schema-tool-with-cache/invoke",
			requestBody: bytes.NewBuffer([]byte(`{}`)),
			wantStatus:  http.StatusOK,
			validateFunc: func(t *testing.T, body string) {
				var result map[string]any
				if err := json.Unmarshal([]byte(body), &result); err != nil {
					t.Fatalf("failed to unmarshal schema result: %v", err)
				}
				// Also check the structure of the schema response for the cached tool.
				expectedKeys := []string{"nodeLabels", "relationships", "constraints", "indexes", "databaseInfo", "statistics"}
				for _, key := range expectedKeys {
					if _, ok := result[key]; !ok {
						t.Errorf("expected key %q not found in schema response", key)
					}
				}
			},
		},
		{
			name:        "invoke my-schema-tool with populated data",
			api:         "http://127.0.0.1:5000/api/tool/my-populated-schema-tool/invoke",
			requestBody: bytes.NewBuffer([]byte(`{}`)),
			wantStatus:  http.StatusOK,
			prepareData: func(t *testing.T) {
				ctx := context.Background()
				driver, err := neo4j.NewDriverWithContext(Neo4jUri, neo4j.BasicAuth(Neo4jUser, Neo4jPass, ""))
				if err != nil {
					t.Fatalf("failed to create neo4j driver: %v", err)
				}

				// Helper to execute queries for setup and teardown.
				execute := func(query string) {
					session := driver.NewSession(ctx, neo4j.SessionConfig{DatabaseName: Neo4jDatabase})
					defer session.Close(ctx)
					// Use ExecuteWrite to ensure the query is committed before proceeding.
					_, err := session.ExecuteWrite(ctx, func(tx neo4j.ManagedTransaction) (any, error) {
						_, err := tx.Run(ctx, query, nil)
						return nil, err
					})

					// Don't fail the test on teardown errors (e.g., entity doesn't exist).
					if err != nil && !strings.Contains(query, "DROP") {
						t.Fatalf("query failed: %s\nerror: %v", query, err)
					}
				}

				// Teardown logic is deferred to ensure it runs even if the test fails.
				// The driver will be closed at the end of this block.
				t.Cleanup(func() {
					execute("DROP CONSTRAINT PersonNameUnique IF EXISTS")
					execute("DROP INDEX MovieTitleIndex IF EXISTS")
					execute("MATCH (n) DETACH DELETE n")
					if err := driver.Close(ctx); err != nil {
						t.Errorf("failed to close driver during cleanup: %v", err)
					}
				})

				// Setup: Create constraints, indexes, and data.
				execute("MERGE (p:Person {name: 'Alice'}) MERGE (m:Movie {title: 'The Matrix'}) MERGE (p)-[:ACTED_IN]->(m)")
				execute("CREATE CONSTRAINT PersonNameUnique IF NOT EXISTS FOR (p:Person) REQUIRE p.name IS UNIQUE")
				execute("CREATE INDEX MovieTitleIndex IF NOT EXISTS FOR (m:Movie) ON (m.title)")
			},
			validateFunc: func(t *testing.T, body string) {
				// Define structs for unmarshaling the detailed schema.
				type Property struct {
					Name  string   `json:"name"`
					Types []string `json:"types"`
				}
				type NodeLabel struct {
					Name       string     `json:"name"`
					Properties []Property `json:"properties"`
				}
				type Relationship struct {
					Type      string `json:"type"`
					StartNode string `json:"startNode"`
					EndNode   string `json:"endNode"`
				}
				type Constraint struct {
					Name       string   `json:"name"`
					Label      string   `json:"label"`
					Properties []string `json:"properties"`
				}
				type Index struct {
					Name       string   `json:"name"`
					Label      string   `json:"label"`
					Properties []string `json:"properties"`
				}
				type Schema struct {
					NodeLabels    []NodeLabel    `json:"nodeLabels"`
					Relationships []Relationship `json:"relationships"`
					Constraints   []Constraint   `json:"constraints"`
					Indexes       []Index        `json:"indexes"`
				}

				var schema Schema
				if err := json.Unmarshal([]byte(body), &schema); err != nil {
					t.Fatalf("failed to unmarshal schema json: %v\nResponse body: %s", err, body)
				}

				// --- Validate Node Labels and Properties ---
				var personLabelFound, movieLabelFound bool
				for _, l := range schema.NodeLabels {
					if l.Name == "Person" {
						personLabelFound = true
						propFound := false
						for _, p := range l.Properties {
							if p.Name == "name" {
								propFound = true
								break
							}
						}
						if !propFound {
							t.Errorf("expected Person label to have 'name' property, but it was not found")
						}
					}
					if l.Name == "Movie" {
						movieLabelFound = true
						propFound := false
						for _, p := range l.Properties {
							if p.Name == "title" {
								propFound = true
								break
							}
						}
						if !propFound {
							t.Errorf("expected Movie label to have 'title' property, but it was not found")
						}
					}
				}
				if !personLabelFound {
					t.Error("expected to find 'Person' in nodeLabels")
				}
				if !movieLabelFound {
					t.Error("expected to find 'Movie' in nodeLabels")
				}

				// --- Validate Relationships ---
				relFound := false
				for _, r := range schema.Relationships {
					if r.Type == "ACTED_IN" && r.StartNode == "Person" && r.EndNode == "Movie" {
						relFound = true
						break
					}
				}
				if !relFound {
					t.Errorf("expected to find relationship '(:Person)-[:ACTED_IN]->(:Movie)', but it was not found")
				}

				// --- Validate Constraints ---
				constraintFound := false
				for _, c := range schema.Constraints {
					if c.Name == "PersonNameUnique" && c.Label == "Person" {
						propFound := false
						for _, p := range c.Properties {
							if p == "name" {
								propFound = true
								break
							}
						}
						if propFound {
							constraintFound = true
							break
						}
					}
				}
				if !constraintFound {
					t.Errorf("expected to find constraint 'PersonNameUnique' on Person(name), but it was not found")
				}

				// --- Validate Indexes ---
				indexFound := false
				for _, i := range schema.Indexes {
					if i.Name == "MovieTitleIndex" && i.Label == "Movie" {
						propFound := false
						for _, p := range i.Properties {
							if p == "title" {
								propFound = true
								break
							}
						}
						if propFound {
							indexFound = true
							break
						}
					}
				}
				if !indexFound {
					t.Errorf("expected to find index 'MovieTitleIndex' on Movie(title), but it was not found")
				}
			},
		},
	}
	for _, tc := range invokeTcs {
		t.Run(tc.name, func(t *testing.T) {
			// Prepare data if a preparation function is provided.
			if tc.prepareData != nil {
				tc.prepareData(t)
			}

			resp, err := http.Post(tc.api, "application/json", tc.requestBody)
			if err != nil {
				t.Fatalf("error when sending a request: %s", err)
			}
			defer resp.Body.Close()
			if resp.StatusCode != tc.wantStatus {
				bodyBytes, _ := io.ReadAll(resp.Body)
				t.Fatalf("response status code: got %d, want %d: %s", resp.StatusCode, tc.wantStatus, string(bodyBytes))
			}

			if tc.wantStatus == http.StatusOK {
				var body map[string]interface{}
				err = json.NewDecoder(resp.Body).Decode(&body)
				if err != nil {
					t.Fatalf("error parsing response body")
				}
				got, ok := body["result"].(string)
				if !ok {
					t.Fatalf("unable to find result in response body")
				}

				if tc.validateFunc != nil {
					// Use the custom validation function if provided.
					tc.validateFunc(t, got)
				} else if got != tc.want {
					// Otherwise, perform a direct string comparison.
					t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
				}
			} else {
				bodyBytes, err := io.ReadAll(resp.Body)
				if err != nil {
					t.Fatalf("failed to read error response body: %s", err)
				}
				bodyString := string(bodyBytes)
				if !strings.Contains(bodyString, tc.wantErrorSubstring) {
					t.Fatalf("response body %q does not contain expected error %q", bodyString, tc.wantErrorSubstring)
				}
			}
		})
	}
}
