// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 main

import (
	"context"
	"os"
	"path/filepath"
	"testing"

	"github.com/apache/beam/sdks/v2/go/container/tools"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/artifact"
	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
	"google.golang.org/protobuf/proto"
)

func TestEnsureEndpointsSet_AllSet(t *testing.T) {
	provisionInfo := &fnpb.ProvisionInfo{
		LoggingEndpoint:  &pipepb.ApiServiceDescriptor{Url: "testLoggingEndpointUrl"},
		ArtifactEndpoint: &pipepb.ApiServiceDescriptor{Url: "testArtifactEndpointUrl"},
		ControlEndpoint:  &pipepb.ApiServiceDescriptor{Url: "testControlEndpointUrl"},
	}
	*loggingEndpoint = ""
	*artifactEndpoint = ""
	*controlEndpoint = ""
	err := ensureEndpointsSet(provisionInfo)
	if err != nil {
		t.Fatalf("ensureEndpointsSet() = %q, want nil", err)
	}
	if got, want := *loggingEndpoint, "testLoggingEndpointUrl"; got != want {
		t.Fatalf("After ensureEndpointsSet(), *loggingEndpoint = %q, want %q", got, want)
	}
	if got, want := *artifactEndpoint, "testArtifactEndpointUrl"; got != want {
		t.Fatalf("After ensureEndpointsSet(), *artifactEndpoint = %q, want %q", got, want)
	}
	if got, want := *controlEndpoint, "testControlEndpointUrl"; got != want {
		t.Fatalf("After ensureEndpointsSet(), *controlEndpoint = %q, want %q", got, want)
	}
}

func TestEnsureEndpointsSet_OneMissing(t *testing.T) {
	provisionInfo := &fnpb.ProvisionInfo{
		LoggingEndpoint:  &pipepb.ApiServiceDescriptor{Url: "testLoggingEndpointUrl"},
		ArtifactEndpoint: &pipepb.ApiServiceDescriptor{Url: "testArtifactEndpointUrl"},
		ControlEndpoint:  &pipepb.ApiServiceDescriptor{Url: ""},
	}
	*loggingEndpoint = ""
	*artifactEndpoint = ""
	*controlEndpoint = ""
	err := ensureEndpointsSet(provisionInfo)
	if err == nil {
		t.Fatalf("ensureEndpointsSet() = nil, want non-nil error")
	}
	if got, want := *loggingEndpoint, "testLoggingEndpointUrl"; got != want {
		t.Fatalf("After ensureEndpointsSet(), *loggingEndpoint = %q, want %q", got, want)
	}
	if got, want := *artifactEndpoint, "testArtifactEndpointUrl"; got != want {
		t.Fatalf("After ensureEndpointsSet(), *artifactEndpoint = %q, want %q", got, want)
	}
	if got, want := *controlEndpoint, ""; got != want {
		t.Fatalf("After ensureEndpointsSet(), *controlEndpoint = %q, want %q", got, want)
	}
}

func TestGetGoWorkerArtifactName_NoArtifacts(t *testing.T) {
	_, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, []*pipepb.ArtifactInformation{})
	if err == nil {
		t.Fatalf("getGoWorkerArtifactName() = nil, want non-nil error")
	}
}

func TestGetGoWorkerArtifactName_OneArtifact(t *testing.T) {
	artifact := constructArtifactInformation(t, artifact.URNGoWorkerBinaryRole, "test/path", "sha")
	artifacts := []*pipepb.ArtifactInformation{&artifact}

	val, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, artifacts)
	if err != nil {
		t.Fatalf("getGoWorkerArtifactName() = %v, want nil", err)
	}
	if got, want := val, "test/path"; got != want {
		t.Fatalf("getGoWorkerArtifactName() = %v, want %v", got, want)
	}
}

func TestGetGoWorkerArtifactName_MultipleArtifactsFirstIsWorker(t *testing.T) {
	artifact1 := constructArtifactInformation(t, artifact.URNGoWorkerBinaryRole, "test/path", "sha")
	artifact2 := constructArtifactInformation(t, "other role", "test/path2", "sha")
	artifacts := []*pipepb.ArtifactInformation{&artifact1, &artifact2}

	val, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, artifacts)
	if err != nil {
		t.Fatalf("getGoWorkerArtifactName() = %v, want nil", err)
	}
	if got, want := val, "test/path"; got != want {
		t.Fatalf("getGoWorkerArtifactName() = %v, want %v", got, want)
	}
}

func TestGetGoWorkerArtifactName_MultipleArtifactsSecondIsWorker(t *testing.T) {
	artifact1 := constructArtifactInformation(t, "other role", "test/path", "sha")
	artifact2 := constructArtifactInformation(t, artifact.URNGoWorkerBinaryRole, "test/path2", "sha")
	artifacts := []*pipepb.ArtifactInformation{&artifact1, &artifact2}

	val, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, artifacts)
	if err != nil {
		t.Fatalf("getGoWorkerArtifactName() = %v, want nil", err)
	}
	if got, want := val, "test/path2"; got != want {
		t.Fatalf("getGoWorkerArtifactName() = %v, want %v", got, want)
	}
}

func TestGetGoWorkerArtifactName_MultipleArtifactsLegacyWay(t *testing.T) {
	artifact1 := constructArtifactInformation(t, "other role", "test/path", "sha")
	artifact2 := constructArtifactInformation(t, "other role", "worker", "sha")
	artifacts := []*pipepb.ArtifactInformation{&artifact1, &artifact2}

	val, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, artifacts)
	if err != nil {
		t.Fatalf("getGoWorkerArtifactName() = %v, want nil", err)
	}
	if got, want := val, "worker"; got != want {
		t.Fatalf("getGoWorkerArtifactName() = %v, want %v", got, want)
	}
}

func TestGetGoWorkerArtifactName_MultipleArtifactsNoneMatch(t *testing.T) {
	artifact1 := constructArtifactInformation(t, "other role", "test/path", "sha")
	artifact2 := constructArtifactInformation(t, "other role", "test/path2", "sha")
	artifacts := []*pipepb.ArtifactInformation{&artifact1, &artifact2}

	_, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, artifacts)
	if err == nil {
		t.Fatalf("getGoWorkerArtifactName() = nil, want non-nil error")
	}
}

func TestCopyExe(t *testing.T) {
	testExeContent := []byte("testContent")

	// Make temp directory to cleanup at the end
	d, err := os.MkdirTemp(os.Getenv("TEST_TMPDIR"), "copyExe-*")
	if err != nil {
		t.Fatalf("failed to make temp directory, got %v", err)
	}
	t.Cleanup(func() { os.RemoveAll(d) })

	// Make our source file and write to it
	src, err := os.CreateTemp(d, "src.exe")
	if err != nil {
		t.Fatalf("failed to make temp file, got %v", err)
	}
	if _, err := src.Write(testExeContent); err != nil {
		t.Fatalf("failed to write to temp file, got %v", err)
	}
	if err := src.Close(); err != nil {
		t.Fatalf("failed to close temp file, got %v", err)
	}

	// Make sure our destination path doesn't exist already
	srcPath, destPath := src.Name(), filepath.Join(d, "dest.exe")
	if _, err := os.Stat(destPath); err == nil {
		t.Fatalf("dest file %v already exists", destPath)
	}

	err = copyExe(srcPath, destPath)
	if err != nil {
		t.Fatalf("copyExe() = %v, want nil", err)
	}
	if _, err := os.Stat(destPath); err != nil {
		t.Fatalf("After running copyExe, os.Stat() = %v, want nil", err)
	}
	destContents, err := os.ReadFile(destPath)
	if err != nil {
		t.Fatalf("After running copyExe, os.ReadFile() = %v, want nil", err)
	}
	if got, want := string(destContents), string(testExeContent); got != want {
		t.Fatalf("After running copyExe, os.ReadFile() = %v, want %v", got, want)
	}
}

func constructArtifactInformation(t *testing.T, roleUrn string, path string, sha string) pipepb.ArtifactInformation {
	t.Helper()

	typePayload, _ := proto.Marshal(&pipepb.ArtifactFilePayload{Path: path, Sha256: sha})

	return pipepb.ArtifactInformation{
		RoleUrn:     roleUrn,
		TypeUrn:     artifact.URNFileArtifact,
		TypePayload: typePayload,
	}
}

func clearEnvVars() {
	_ = os.Unsetenv(cloudProfilingJobName)
	_ = os.Unsetenv(cloudProfilingJobID)
}

func TestConfigureGoogleCloudProfilerEnvVars(t *testing.T) {
	tests := []struct {
		name           string
		options        string
		metadata       map[string]string
		expectedName   string
		expectedID     string
		expectingError bool
	}{
		{
			name: "Profiler name from options",
			options: `{
				"beam:option:go_options:v1": {
					"options": {
						"dataflow_service_options": "enable_google_cloud_profiler=custom_profiler"
					}
				}
			}`,
			metadata: map[string]string{
				"job_id": "job-123",
			},
			expectedName:   "custom_profiler",
			expectedID:     "job-123",
			expectingError: false,
		},
		{
			name: "Fallback to job_name",
			options: `{
				"beam:option:go_options:v1": {
					"options": {
					    "dataflow_service_options": "enable_google_cloud_profiler"
					}
				}
			}`,
			metadata: map[string]string{
				"job_name": "fallback_profiler",
				"job_id":   "job-456",
			},
			expectedName:   "fallback_profiler",
			expectedID:     "job-456",
			expectingError: false,
		},
		{
			name: "Missing job_id",
			options: `{
				"beam:option:go_options:v1": {
					"options": {
						"dataflow_service_options": "enable_google_cloud_profiler=custom_profiler"
					}
				}
			}`,
			metadata: map[string]string{
				"job_name": "custom_profiler",
			},
			expectingError: true,
		},
		{
	        name: "Missing profiler name and job_name",
	        options: `{
		        "beam:option:go_options:v1": {
			        "options": {
				        "dataflow_service_options": "enable_google_cloud_profiler"
			        }
		        }
	        }`,
	        metadata: map[string]string{
		        "job_id": "job-789",
	        },
	        expectingError: true,
        },
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			clearEnvVars()
			ctx := context.Background()

			err := configureGoogleCloudProfilerEnvVars(ctx, &tools.Logger{}, tt.metadata, tt.options)

			if tt.expectingError {
				if err == nil {
					t.Errorf("Expected error but got nil")
				}
				return
			} else {
				if err != nil {
					t.Errorf("Did not expect error but got: %v", err)
					return
				}
			}

			gotName := os.Getenv(cloudProfilingJobName)
			gotID := os.Getenv(cloudProfilingJobID)

			if gotName != tt.expectedName {
				t.Errorf("Expected profiler name '%s', got '%s'", tt.expectedName, gotName)
			}
			if gotID != tt.expectedID {
				t.Errorf("Expected job ID '%s', got '%s'", tt.expectedID, gotID)
			}
		})
	}
}